Spaces:
Running on Zero
Running on Zero
File size: 4,122 Bytes
3e33e2e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | """Extract historical bars for key liquid symbols into SQLite and JSON cache."""
from __future__ import annotations
import argparse
import json
import sqlite3
import time
from pathlib import Path
from drift_quant.nse_data.loader import (
load_selected_symbols,
save_selected_cache,
)
CORE_UNIVERSE = (
# Banking
"HDFCBANK",
"ICICIBANK",
"SBIN",
"AXISBANK",
"KOTAKBANK",
"INDUSINDBK",
"BANKBARODA",
# Financial Services
"BAJFINANCE",
"BAJAJFINSV",
"HDFCLIFE",
"SBILIFE",
# Information Technology
"TCS",
"INFY",
"WIPRO",
"HCLTECH",
"TECHM",
"LTIM",
# Energy, Oil & Gas
"RELIANCE",
"ONGC",
"BPCL",
"POWERGRID",
"NTPC",
"TATAPOWER",
"COALINDIA",
# FMCG & Consumer
"ITC",
"HINDUNILVR",
"NESTLEIND",
"BRITANNIA",
"DABUR",
"TATACONSUM",
# Automobile
"TATAMOTORS",
"MARUTI",
"M&M",
"BAJAJ-AUTO",
"EICHERMOT",
"HEROMOTOCO",
# Metals & Mining
"TATASTEEL",
"JSWSTEEL",
"HINDALCO",
"VEDL",
"JINDALSTEL",
# Pharma & Healthcare
"SUNPHARMA",
"DRREDDY",
"CIPLA",
"DIVISLAB",
"APOLLOHOSP",
# Infrastructure, Industrials & Conglomerates
"LT",
"BHARTIARTL",
"ULTRACEMCO",
"GRASIM",
"ADANIENT",
"ADANIPORTS",
"TITAN",
"ASIANPAINT",
)
def populate_sqlite(dataset, db_path: Path) -> None:
db_path.parent.mkdir(parents=True, exist_ok=True)
if db_path.is_file():
db_path.unlink()
conn = sqlite3.connect(db_path)
cur = conn.cursor()
cur.execute(
"""
CREATE TABLE bars (
symbol TEXT NOT NULL,
timestamp TEXT NOT NULL,
open REAL NOT NULL,
high REAL NOT NULL,
low REAL NOT NULL,
close REAL NOT NULL,
volume REAL NOT NULL,
source_row INTEGER NOT NULL
)
"""
)
records = []
for _symbol, bars in dataset.bars_by_symbol.items():
for bar in bars:
records.append(
(
bar.symbol,
bar.timestamp.isoformat(sep=" "),
bar.open,
bar.high,
bar.low,
bar.close,
bar.volume,
bar.source_row,
)
)
cur.executemany(
"""
INSERT INTO bars (symbol, timestamp, open, high, low, close, volume, source_row)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""",
records,
)
cur.execute("CREATE INDEX idx_bars_symbol_ts ON bars (symbol, timestamp)")
conn.commit()
conn.close()
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"archive",
type=Path,
nargs="?",
default=Path("all_NSE500_15min_data_2019_onwards.zip"),
)
parser.add_argument("--audit", type=Path, default=Path("artifacts/dataset_audit.json"))
parser.add_argument(
"--json-cache", type=Path, default=Path("artifacts/cache/nse_universe_20.json")
)
parser.add_argument(
"--sqlite-cache", type=Path, default=Path("artifacts/cache/market_data.sqlite")
)
args = parser.parse_args()
audit = json.loads(args.audit.read_text(encoding="utf-8"))
sha = audit["archive_sha256"]
print(f"Scanning {args.archive} for {len(CORE_UNIVERSE)} symbols across 2019-2025...")
t0 = time.time()
dataset = load_selected_symbols(
args.archive,
sha,
CORE_UNIVERSE,
progress=print,
)
elapsed = time.time() - t0
print(
f"Extracted {dataset.rows_selected:,} bars across {len(dataset.symbols)} symbols "
f"in {elapsed:.1f}s."
)
print(f"Writing JSON cache to {args.json_cache}...")
save_selected_cache(dataset, args.json_cache)
print(f"Writing SQLite database to {args.sqlite_cache}...")
populate_sqlite(dataset, args.sqlite_cache)
print("Market data caching complete.")
if __name__ == "__main__":
main()
|