Spaces:
Running on Zero
Running on Zero
File size: 7,017 Bytes
83381c9 | 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 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 | """Incremental loading of selected NSE symbols directly from the source ZIP."""
from __future__ import annotations
import csv
import json
import zipfile
from collections.abc import Callable, Iterable
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path
from drift_quant.nse_data.audit import (
EXPECTED_COLUMNS,
TIMESTAMP_FORMAT,
ArchiveAuditError,
_is_valid_ohlcv,
_select_csv_member,
)
@dataclass(frozen=True, slots=True)
class NseBar:
"""A raw NSE bar preserving its source timestamp and provenance."""
timestamp: datetime
symbol: str
open: float
high: float
low: float
close: float
volume: float
source_row: int
@dataclass(frozen=True, slots=True)
class LoadedNseDataset:
"""Selected raw observations and explicit exclusion accounting."""
archive_path: str
archive_sha256: str
symbols: tuple[str, ...]
bars_by_symbol: dict[str, tuple[NseBar, ...]]
rows_scanned: int
rows_selected: int
invalid_rows_excluded: int
@property
def dataset_id(self) -> str:
selection = ",".join(self.symbols)
return f"nse15m:{self.archive_sha256[:16]}:{selection}"
def _parse_selected_row(row: dict[str, str | None], row_number: int) -> NseBar | None:
try:
timestamp = datetime.strptime(row["date"] or "", TIMESTAMP_FORMAT)
open_price = float(row["open"] or "")
high = float(row["high"] or "")
low = float(row["low"] or "")
close = float(row["close"] or "")
volume = float(row["volume"] or "")
symbol = (row["symbol"] or "").strip().upper()
except (KeyError, TypeError, ValueError):
return None
if not symbol or not _is_valid_ohlcv(open_price, high, low, close, volume):
return None
return NseBar(
timestamp=timestamp,
symbol=symbol,
open=open_price,
high=high,
low=low,
close=close,
volume=volume,
source_row=row_number,
)
def load_selected_symbols(
archive_path: Path,
archive_sha256: str,
symbols: Iterable[str],
*,
start: datetime | None = None,
end: datetime | None = None,
progress: Callable[[str], None] | None = None,
) -> LoadedNseDataset:
"""Read only selected symbols with one streaming pass over the ZIP member.
`start` is inclusive and `end` is exclusive. Invalid source observations are
counted and excluded; no prices are filled or repaired.
"""
requested = tuple(sorted({symbol.strip().upper() for symbol in symbols if symbol.strip()}))
if not requested:
raise ArchiveAuditError("At least one symbol is required.")
if start is not None and end is not None and start >= end:
raise ArchiveAuditError("Start timestamp must precede end timestamp.")
if not archive_path.is_file():
raise ArchiveAuditError(f"Archive does not exist: {archive_path}")
requested_set = set(requested)
bars: dict[str, list[NseBar]] = {symbol: [] for symbol in requested}
rows_scanned = 0
rows_selected = 0
invalid_rows_excluded = 0
with zipfile.ZipFile(archive_path) as archive:
member = _select_csv_member(archive)
with archive.open(member) as binary_stream:
from io import TextIOWrapper
with TextIOWrapper(binary_stream, encoding="utf-8", newline="") as stream:
reader = csv.DictReader(stream)
if tuple(reader.fieldnames or ()) != EXPECTED_COLUMNS:
raise ArchiveAuditError(f"Unexpected CSV schema: {reader.fieldnames!r}")
for row in reader:
rows_scanned += 1
raw_symbol = (row.get("symbol") or "").strip().upper()
if raw_symbol not in requested_set:
continue
bar = _parse_selected_row(row, rows_scanned)
if bar is None:
invalid_rows_excluded += 1
continue
if (start is not None and bar.timestamp < start) or (
end is not None and bar.timestamp >= end
):
continue
bars[bar.symbol].append(bar)
rows_selected += 1
if progress is not None and rows_scanned % 1_000_000 == 0:
progress(
f"Scanned {rows_scanned:,} rows; selected "
f"{rows_selected:,} valid observations."
)
normalized = {
symbol: tuple(sorted(symbol_bars, key=lambda bar: bar.timestamp))
for symbol, symbol_bars in bars.items()
}
return LoadedNseDataset(
archive_path=str(archive_path.resolve()),
archive_sha256=archive_sha256,
symbols=requested,
bars_by_symbol=normalized,
rows_scanned=rows_scanned,
rows_selected=rows_selected,
invalid_rows_excluded=invalid_rows_excluded,
)
def save_selected_cache(dataset: LoadedNseDataset, output_path: Path) -> None:
"""Persist a compact, explicit cache for repeat runs over the same selection."""
output_path.parent.mkdir(parents=True, exist_ok=True)
serializable = {
"archive_path": dataset.archive_path,
"archive_sha256": dataset.archive_sha256,
"symbols": dataset.symbols,
"rows_scanned": dataset.rows_scanned,
"rows_selected": dataset.rows_selected,
"invalid_rows_excluded": dataset.invalid_rows_excluded,
"bars_by_symbol": {
symbol: [{**asdict(bar), "timestamp": bar.timestamp.isoformat(sep=" ")} for bar in bars]
for symbol, bars in dataset.bars_by_symbol.items()
},
}
output_path.write_text(json.dumps(serializable, separators=(",", ":")), encoding="utf-8")
def load_selected_cache(cache_path: Path) -> LoadedNseDataset:
"""Load an explicit cache created by :func:`save_selected_cache`."""
payload = json.loads(cache_path.read_text(encoding="utf-8"))
return LoadedNseDataset(
archive_path=payload["archive_path"],
archive_sha256=payload["archive_sha256"],
symbols=tuple(payload["symbols"]),
bars_by_symbol={
symbol: tuple(
NseBar(
timestamp=datetime.fromisoformat(bar["timestamp"]),
symbol=bar["symbol"],
open=bar["open"],
high=bar["high"],
low=bar["low"],
close=bar["close"],
volume=bar["volume"],
source_row=bar["source_row"],
)
for bar in bars
)
for symbol, bars in payload["bars_by_symbol"].items()
},
rows_scanned=payload["rows_scanned"],
rows_selected=payload["rows_selected"],
invalid_rows_excluded=payload["invalid_rows_excluded"],
)
|