Spaces:
Running on Zero
Running on Zero
| """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, | |
| ) | |
| 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 | |
| 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 | |
| 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"], | |
| ) | |