SPB Git

spb/findata

Public
11.5 KB · 222 lines python
1"""Zip archive handling and CSV → canonical Arrow conversion for FirstRate bar files.23Vendor file facts: `{TICKER}_{period}_{timeframe}_{adjustment}.txt` (update archives) or `{TICKER}_{timeframe}.txt`4(full archives), no header, `yyyy-MM-dd HH:mm:ss,O,H,L,C,V` (daily: `yyyy-MM-dd,O,H,L,C,V[,OI]`), FX intraday splits5date and time into two fields, line endings may mix LF/CRLF, nested zips occur in bundle archives.6"""7from __future__ import annotations89import io10import re11import zipfile12from collections.abc import Iterator13from datetime import datetime14from pathlib import Path1516import polars as pl17import pyarrow as pa18from findata_schema import BARS_SCHEMA, Adjustment, AssetClass, Interval1920VENDOR_TF = {"1min": Interval.m1, "5min": Interval.m5, "30min": Interval.m30, "1hour": Interval.h1, "1day": Interval.d1}21VENDOR_ADJ = {"UNADJUSTED": Adjustment.raw, "adj_split": Adjustment.split, "adj_splitdiv": Adjustment.split_div}22_TICKER_RE = re.compile(r"^([A-Za-z0-9.\-]+?)(?:[_-](?:full|day|week|month|1min|5min|30min|1hour|1day|UNADJUSTED|adj_split|adj_splitdiv|contin.*))*\.(txt|csv)$", re.IGNORECASE)232425_TF_SUFFIX_RE = re.compile(r"_(1min|5min|30min|1hour|1day)$", re.I)262728def ticker_from_filename(name: str, contract: bool = False) -> str:29    base = Path(name).name30    if contract:  # individual futures contracts: 'A6_F25_1day.txt' -> 'A6_F25'31        stem = Path(base).stem32        return _TF_SUFFIX_RE.sub("", stem).upper()33    m = _TICKER_RE.match(base)34    if m:35        return m.group(1).upper()36    return base.split("_")[0].upper()373839def iter_archive_files(zip_path: Path) -> Iterator[tuple[str, bytes]]:40    """Yield (member_name, bytes) for every data file, descending into nested zips.4142    Some vendor archives (option chains from 2023 on) use Deflate64, which Python's zipfile cannot decompress43    (`NotImplementedError: That compression method is not supported`). In that case the archive is extracted with44    the system `7z` into a temporary directory next to the archive and the files are streamed from disk."""45    try:46        yield from _iter_zip_members(zip_path)47    except NotImplementedError:48        yield from _iter_via_7z(zip_path)495051def _iter_zip_members(zip_path: Path) -> Iterator[tuple[str, bytes]]:52    with zipfile.ZipFile(zip_path) as zf:53        for info in zf.infolist():54            if info.is_dir():55                continue56            name = info.filename57            if name.lower().endswith(".zip"):58                data = zf.read(info)59                with zipfile.ZipFile(io.BytesIO(data)) as inner:60                    for inner_info in inner.infolist():61                        if not inner_info.is_dir() and inner_info.filename.lower().endswith((".txt", ".csv")):62                            yield f"{name}!/{inner_info.filename}", inner.read(inner_info)63            elif name.lower().endswith((".txt", ".csv")):64                yield name, zf.read(info)656667def _find_7z() -> str | None:68    import shutil69    for cand in ("7zz", "7z", "7za"):70        p = shutil.which(cand) or shutil.which(cand, path="/opt/homebrew/bin:/usr/local/bin")71        if p:72            return p73    return None747576def _iter_via_7z(zip_path: Path) -> Iterator[tuple[str, bytes]]:77    import shutil78    import subprocess79    import tempfile80    tool = _find_7z()81    if not tool:82        raise RuntimeError(f"{zip_path.name}: unsupported zip compression (Deflate64) and no 7z binary found (brew install p7zip)")83    tmp = Path(tempfile.mkdtemp(prefix="unz-", dir=str(zip_path.parent)))84    try:85        r = subprocess.run([tool, "x", "-y", "-bd", "-bso0", "-bsp0", f"-o{tmp}", str(zip_path)], capture_output=True, text=True, timeout=6 * 3600)86        if r.returncode not in (0, 1):  # 1 = warnings87            raise RuntimeError(f"7z extraction failed for {zip_path.name}: {r.stderr[-300:]}")88        for p in sorted(tmp.rglob("*")):89            if p.is_file() and p.suffix.lower() in (".txt", ".csv"):90                yield str(p.relative_to(tmp)), p.read_bytes()91            elif p.is_file() and p.suffix.lower() == ".zip":92                yield from ((f"{p.name}!/{n}", b) for n, b in iter_archive_files(p))93    finally:94        shutil.rmtree(tmp, ignore_errors=True)959697def parse_bar_file(data: bytes, symbol: str, asset_class: AssetClass, interval: Interval, adjustment: Adjustment,98                   source_id: str, ingested_at: datetime, source_tz: str = "America/New_York",99                   currency: str = "USD", venue: str | None = None, revision_id: int = 1) -> tuple[pa.Table, int]:100    """Parse one vendor CSV into the canonical bars table. Returns (table, raw_row_count).101102    Duplicate (symbol, timestamp) rows are *kept* here and flagged downstream; we never drop source records silently.103    Unparseable timestamps (e.g. non-existent local time at DST spring-forward) are dropped from the typed table and104    counted in `raw_row_count - table.num_rows`.105    """106    text = data.decode("utf-8", errors="replace").replace("\r\n", "\n").replace("\r", "\n")107    lines = [ln for ln in text.split("\n") if ln.strip()]108    if lines and lines[0].lower().startswith(("datetime", "date", "timestamp")):109        lines = lines[1:]110    raw_rows = len(lines)111    if raw_rows == 0:112        return pa.table({f.name: pa.array([], type=f.type) for f in BARS_SCHEMA}, schema=BARS_SCHEMA), 0113    # field count = the most common count over a sample (vendor files occasionally contain ragged lines)114    from collections import Counter115    ncols = Counter(len(ln.split(",")) for ln in lines[:500]).most_common(1)[0][0]116    buf = "\n".join(lines).encode()117    if ncols == 7 and re.match(r"^\d{8},\d{2}:\d{2}", lines[0]):118        names = ["d", "t", "open", "high", "low", "close", "volume"]119        df = pl.read_csv(buf, has_header=False, new_columns=names, infer_schema_length=0, truncate_ragged_lines=True)120        df = df.with_columns((pl.col("d") + " " + pl.col("t")).alias("dt_raw")).drop("d", "t")121        fmt = "%Y%m%d %H:%M:%S"122    else:123        if ncols == 5:  # indices: no volume column124            names = ["dt_raw", "open", "high", "low", "close"]125        else:126            names = ["dt_raw", "open", "high", "low", "close", "volume"] + (["open_interest"] if ncols == 7 else [])127        if ncols not in (5, 6, 7):128            raise ValueError(f"unexpected column count {ncols} in {symbol}")129        df = pl.read_csv(buf, has_header=False, new_columns=names, infer_schema_length=0, truncate_ragged_lines=True)130        if ncols == 5:131            df = df.with_columns(pl.lit("0").alias("volume"))132        fmt = "%Y-%m-%d %H:%M:%S" if interval != Interval.d1 else None133    for c in ("open", "high", "low", "close", "volume"):134        df = df.with_columns(pl.col(c).str.strip_chars().cast(pl.Float64, strict=False))135    if interval == Interval.d1:136        local = pl.col("dt_raw").str.strip_chars().str.slice(0, 10).str.to_date("%Y-%m-%d", strict=False).cast(pl.Datetime("us"))137    else:138        local = pl.coalesce(139            pl.col("dt_raw").str.strip_chars().str.to_datetime(fmt, strict=False, time_unit="us"),140            pl.col("dt_raw").str.strip_chars().str.to_datetime("%Y-%m-%d %H:%M", strict=False, time_unit="us"),141        )142    df = df.with_columns(local.alias("local_naive"))143    df = df.filter(pl.col("local_naive").is_not_null() & pl.col("open").is_not_null() & pl.col("close").is_not_null())144    df = df.with_columns(145        pl.col("local_naive").dt.replace_time_zone(source_tz, ambiguous="earliest", non_existent="null").alias("ts_local")146    ).filter(pl.col("ts_local").is_not_null())147    df = df.with_columns(148        pl.col("ts_local").dt.convert_time_zone("UTC").alias("timestamp_utc"),149        pl.col("ts_local").dt.date().alias("session_date"),150    )151    n = df.height152    out = pl.DataFrame({153        "instrument_id": pl.repeat(f"{source_id}:{asset_class.value}:{symbol}", n, eager=True),154        "symbol_as_reported": pl.repeat(symbol, n, eager=True),155        "venue": pl.repeat(venue, n, eager=True, dtype=pl.Utf8),156        "asset_class": pl.repeat(asset_class.value, n, eager=True),157        "interval": pl.repeat(interval.value, n, eager=True),158        "timestamp_utc": df["timestamp_utc"],159        "session_date": df["session_date"],160        "open": df["open"], "high": df["high"], "low": df["low"], "close": df["close"],161        "volume": df["volume"].fill_null(0.0),162        "currency": pl.repeat(currency, n, eager=True),163        "trade_count": pl.repeat(None, n, eager=True, dtype=pl.Int32),164        "vwap": pl.repeat(None, n, eager=True, dtype=pl.Float64),165        "adjustment_type": pl.repeat(adjustment.value, n, eager=True),166        "source_id": pl.repeat(source_id, n, eager=True),167        "ingested_at": pl.repeat(ingested_at, n, eager=True, dtype=pl.Datetime("us", "UTC")),168        "quality_flags": pl.repeat([], n, eager=True, dtype=pl.List(pl.Utf8)),169        "revision_id": pl.repeat(revision_id, n, eager=True, dtype=pl.Int32),170    })171    return out.to_arrow().cast(BARS_SCHEMA), raw_rows172173174def parse_meta_pairs(data: bytes) -> list[tuple[str, float]]:175    """Splits/dividends per-ticker files: `yyyy-MM-dd,value` lines."""176    out = []177    for ln in data.decode("utf-8", errors="replace").splitlines():178        ln = ln.strip()179        if not ln or ln.lower().startswith("date"):180            continue181        parts = ln.split(",")182        if len(parts) >= 2:183            try:184                out.append((parts[0].strip(), float(parts[1])))185            except ValueError:186                continue187    return out188189190def parse_options_file(data: bytes, underlying: str, source_id: str, ingested_at: datetime) -> tuple[pa.Table, int]:191    """Vendor option-chain CSV (16 columns, no header): Trade Date, Strike, Expiry, Call/Put, Last, Bid, Ask, Bid IV,192    Ask IV, Open Interest, Volume, Delta, Gamma, Vega, Theta, Rho → OPTIONS_SCHEMA."""193    from findata_schema import OPTIONS_SCHEMA194    text = data.decode("utf-8", errors="replace").replace("\r\n", "\n").replace("\r", "\n")195    lines = [ln for ln in text.split("\n") if ln.strip()]196    if lines and lines[0].lower().startswith(("trade", "date")):197        lines = lines[1:]198    raw_rows = len(lines)199    empty = pa.table({f.name: pa.array([], type=f.type) for f in OPTIONS_SCHEMA}, schema=OPTIONS_SCHEMA)200    if raw_rows == 0:201        return empty, 0202    names = ["trade_date", "strike", "expiry", "call_put", "last_price", "bid", "ask", "bid_iv", "ask_iv", "open_interest",203             "volume", "delta", "gamma", "vega", "theta", "rho"]204    from collections import Counter205    ncols = Counter(len(ln.split(",")) for ln in lines[:500]).most_common(1)[0][0]206    if ncols != 16:207        raise ValueError(f"unexpected option-chain column count {ncols} for {underlying}")208    df = pl.read_csv("\n".join(lines).encode(), has_header=False, new_columns=names, infer_schema_length=0, truncate_ragged_lines=True)209    num = [c for c in names if c not in ("trade_date", "expiry", "call_put")]210    df = df.with_columns([pl.col(c).str.strip_chars().cast(pl.Float64, strict=False) for c in num])211    df = df.with_columns(212        pl.col("trade_date").str.strip_chars().str.slice(0, 10).str.to_date("%Y-%m-%d", strict=False),213        pl.col("expiry").str.strip_chars().str.slice(0, 10).str.to_date("%Y-%m-%d", strict=False),214        pl.col("call_put").str.strip_chars().str.to_lowercase(),215    ).filter(pl.col("trade_date").is_not_null() & pl.col("expiry").is_not_null() & pl.col("strike").is_not_null())216    n = df.height217    out = df.with_columns(218        pl.lit(underlying).alias("underlying"), pl.lit(source_id).alias("source_id"),219        pl.lit(ingested_at).cast(pl.Datetime("us", "UTC")).alias("ingested_at"),220    ).select(OPTIONS_SCHEMA.names)221    return out.to_arrow().cast(OPTIONS_SCHEMA), raw_rows222
Files · main
Menu

Theme