SPB Git

spb/findata

Public
6.6 KB · 131 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)232425def ticker_from_filename(name: str) -> str:26    base = Path(name).name27    m = _TICKER_RE.match(base)28    if m:29        return m.group(1).upper()30    return base.split("_")[0].upper()313233def iter_archive_files(zip_path: Path) -> Iterator[tuple[str, bytes]]:34    """Yield (member_name, bytes) for every data file, descending into nested zips."""35    with zipfile.ZipFile(zip_path) as zf:36        for info in zf.infolist():37            if info.is_dir():38                continue39            name = info.filename40            if name.lower().endswith(".zip"):41                data = zf.read(info)42                with zipfile.ZipFile(io.BytesIO(data)) as inner:43                    for inner_info in inner.infolist():44                        if not inner_info.is_dir() and inner_info.filename.lower().endswith((".txt", ".csv")):45                            yield f"{name}!/{inner_info.filename}", inner.read(inner_info)46            elif name.lower().endswith((".txt", ".csv")):47                yield name, zf.read(info)484950def parse_bar_file(data: bytes, symbol: str, asset_class: AssetClass, interval: Interval, adjustment: Adjustment,51                   source_id: str, ingested_at: datetime, source_tz: str = "America/New_York",52                   currency: str = "USD", venue: str | None = None, revision_id: int = 1) -> tuple[pa.Table, int]:53    """Parse one vendor CSV into the canonical bars table. Returns (table, raw_row_count).5455    Duplicate (symbol, timestamp) rows are *kept* here and flagged downstream; we never drop source records silently.56    Unparseable timestamps (e.g. non-existent local time at DST spring-forward) are dropped from the typed table and57    counted in `raw_row_count - table.num_rows`.58    """59    text = data.decode("utf-8", errors="replace").replace("\r\n", "\n").replace("\r", "\n")60    lines = [ln for ln in text.split("\n") if ln.strip()]61    if lines and lines[0].lower().startswith(("datetime", "date", "timestamp")):62        lines = lines[1:]63    raw_rows = len(lines)64    if raw_rows == 0:65        return pa.table({f.name: pa.array([], type=f.type) for f in BARS_SCHEMA}, schema=BARS_SCHEMA), 066    ncols = len(lines[0].split(","))67    buf = "\n".join(lines).encode()68    if ncols == 7 and re.match(r"^\d{8},\d{2}:\d{2}", lines[0]):69        names = ["d", "t", "open", "high", "low", "close", "volume"]70        df = pl.read_csv(buf, has_header=False, new_columns=names, infer_schema_length=0)71        df = df.with_columns((pl.col("d") + " " + pl.col("t")).alias("dt_raw")).drop("d", "t")72        fmt = "%Y%m%d %H:%M:%S"73    else:74        names = ["dt_raw", "open", "high", "low", "close", "volume"] + (["open_interest"] if ncols == 7 else [])75        if ncols not in (6, 7):76            raise ValueError(f"unexpected column count {ncols} in {symbol}")77        df = pl.read_csv(buf, has_header=False, new_columns=names, infer_schema_length=0)78        fmt = "%Y-%m-%d %H:%M:%S" if interval != Interval.d1 else None79    for c in ("open", "high", "low", "close", "volume"):80        df = df.with_columns(pl.col(c).str.strip_chars().cast(pl.Float64, strict=False))81    if interval == Interval.d1:82        local = pl.col("dt_raw").str.strip_chars().str.slice(0, 10).str.to_date("%Y-%m-%d", strict=False).cast(pl.Datetime("us"))83    else:84        local = pl.col("dt_raw").str.strip_chars().str.to_datetime(fmt, strict=False, time_unit="us")85    df = df.with_columns(local.alias("local_naive"))86    df = df.filter(pl.col("local_naive").is_not_null() & pl.col("open").is_not_null() & pl.col("close").is_not_null())87    df = df.with_columns(88        pl.col("local_naive").dt.replace_time_zone(source_tz, ambiguous="earliest", non_existent="null").alias("ts_local")89    ).filter(pl.col("ts_local").is_not_null())90    df = df.with_columns(91        pl.col("ts_local").dt.convert_time_zone("UTC").alias("timestamp_utc"),92        pl.col("ts_local").dt.date().alias("session_date"),93    )94    n = df.height95    out = pl.DataFrame({96        "instrument_id": pl.repeat(f"{source_id}:{asset_class.value}:{symbol}", n, eager=True),97        "symbol_as_reported": pl.repeat(symbol, n, eager=True),98        "venue": pl.repeat(venue, n, eager=True, dtype=pl.Utf8),99        "asset_class": pl.repeat(asset_class.value, n, eager=True),100        "interval": pl.repeat(interval.value, n, eager=True),101        "timestamp_utc": df["timestamp_utc"],102        "session_date": df["session_date"],103        "open": df["open"], "high": df["high"], "low": df["low"], "close": df["close"],104        "volume": df["volume"].fill_null(0.0),105        "currency": pl.repeat(currency, n, eager=True),106        "trade_count": pl.repeat(None, n, eager=True, dtype=pl.Int32),107        "vwap": pl.repeat(None, n, eager=True, dtype=pl.Float64),108        "adjustment_type": pl.repeat(adjustment.value, n, eager=True),109        "source_id": pl.repeat(source_id, n, eager=True),110        "ingested_at": pl.repeat(ingested_at, n, eager=True, dtype=pl.Datetime("us", "UTC")),111        "quality_flags": pl.repeat([], n, eager=True, dtype=pl.List(pl.Utf8)),112        "revision_id": pl.repeat(revision_id, n, eager=True, dtype=pl.Int32),113    })114    return out.to_arrow().cast(BARS_SCHEMA), raw_rows115116117def parse_meta_pairs(data: bytes) -> list[tuple[str, float]]:118    """Splits/dividends per-ticker files: `yyyy-MM-dd,value` lines."""119    out = []120    for ln in data.decode("utf-8", errors="replace").splitlines():121        ln = ln.strip()122        if not ln or ln.lower().startswith("date"):123            continue124        parts = ln.split(",")125        if len(parts) >= 2:126            try:127                out.append((parts[0].strip(), float(parts[1])))128            except ValueError:129                continue130    return out131
Files · de3e9e393936b8d9c1340b2c85d7ccb9ffa8264a
Menu

Theme