Greg Pomerantz
9263283995
overrides/corrections/{SYM}.json (git-tracked, with as_of/source/note)
holds remove/replace/add ops on the dividend and capital-gain series.
data.py applies them on top of whatever the data root (or the frozen
snapshot) provides, in both the full build and the incremental refresh
path, and the corrections dir joins the cache manifest so a change
invalidates the cache. A goget re-download of the base CSV can never
clobber a confirmed correction. Format and usage documented in data.py.
497 lines
19 KiB
Python
497 lines
19 KiB
Python
"""Data ingestion: Yahoo Finance CSV dumps -> parquet cache.
|
|
|
|
Layout expected in the data root (default ~/prog/fin/stocks):
|
|
{SYM}-history.csv Date,Open,High,Low,Close,Adj Close,Volume
|
|
{SYM}-dividend.csv Date,Dividends (optional)
|
|
{SYM}-capitalGain.csv Date,Capital Gains (optional)
|
|
{SYM}.json Yahoo chart API payload (metadata only)
|
|
|
|
All panels are date x symbol DataFrames. The "adj" panel (Adj Close)
|
|
already bakes in dividends and capital gains, so it is the correct
|
|
input for pre-tax total returns. The dividend/capgain panels are used
|
|
only by the tax engine (to tax distributions).
|
|
|
|
First call builds a parquet cache (default: .cache/) in ~1 min for
|
|
4k symbols; subsequent calls load in well under a second.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import threading
|
|
import time
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import pandas as pd
|
|
|
|
DEFAULT_ROOT = Path("~/prog/fin/stocks").expanduser()
|
|
CACHE_DIR = Path(__file__).parent / ".cache"
|
|
# Frozen overrides: stale (Yahoo-empty) tickers whose final series is
|
|
# snapshotted here by scripts/audit_stale.py. For symbols that have a
|
|
# frozen copy, the data-root files are ignored so a future goget
|
|
# re-download can never clobber the last known-good series.
|
|
FROZEN_DIR = Path(__file__).parent / "overrides" / "frozen"
|
|
# Corrections: per-symbol distribution fixes confirmed against official
|
|
# sources (see scripts/verify_official.py). Applied on top of whatever the
|
|
# data root / frozen dir says, so re-downloads can never clobber them.
|
|
# Format, one file per symbol in overrides/corrections/:
|
|
# {"dividends": {"remove": [["2008-12-18", 0.292], ...],
|
|
# "replace": {"2008-12-18": 0.0},
|
|
# "add": [["2007-12-18", 0.292], ...]},
|
|
# "capitalGains": {...same ops...},
|
|
# "as_of": "2026-08-31", "source": "<filing URL>", "note": "..."}
|
|
CORRECTIONS_DIR = Path(__file__).parent / "overrides" / "corrections"
|
|
_corr_cache: dict = {}
|
|
|
|
|
|
def _load_corrections(sym: str) -> dict:
|
|
f = CORRECTIONS_DIR / f"{sym.upper()}.json"
|
|
if f.exists():
|
|
try:
|
|
if f not in _corr_cache or _corr_cache[f][0] != f.stat().st_mtime_ns:
|
|
_corr_cache[f] = (f.stat().st_mtime_ns,
|
|
json.loads(f.read_text()))
|
|
return _corr_cache[f][1]
|
|
except Exception:
|
|
return {}
|
|
return {}
|
|
|
|
|
|
def _apply_corrections(sym: str, suffix: str, ser: pd.Series) -> pd.Series:
|
|
key = {"dividend": "dividends", "capitalGain": "capitalGains"}.get(suffix)
|
|
ops = _load_corrections(sym).get(key) if key else None
|
|
if not ops:
|
|
return ser
|
|
d = {ts: float(v) for ts, v in ser.items()}
|
|
for r in ops.get("remove", []):
|
|
ts = pd.Timestamp(r[0])
|
|
amt = r[1] if len(r) > 1 else None
|
|
if ts in d and (amt is None or abs(d[ts] - amt) < 1e-9):
|
|
del d[ts]
|
|
for ts_s, amt in ops.get("replace", {}).items():
|
|
d[pd.Timestamp(ts_s)] = float(amt)
|
|
for r in ops.get("add", []):
|
|
ts = pd.Timestamp(r[0])
|
|
d[ts] = d.get(ts, 0.0) + float(r[1])
|
|
if not d:
|
|
return pd.Series(dtype=float)
|
|
return pd.Series(d).sort_index()
|
|
|
|
_SUFFIXES = ("history", "dividend", "capitalGain")
|
|
_PANELS = ("adj", "close", "div", "capg")
|
|
|
|
|
|
@dataclass
|
|
class Bundle:
|
|
adj: pd.DataFrame # adjusted close (total-return adjusted)
|
|
close: pd.DataFrame # raw close
|
|
div: pd.DataFrame # dividend distributions per share
|
|
capg: pd.DataFrame # capital gain distributions per share
|
|
names: dict # symbol -> long name (best effort)
|
|
|
|
|
|
def _read_csv_clean(f: Path) -> pd.DataFrame:
|
|
df = pd.read_csv(f, index_col=0)
|
|
df.index = pd.to_datetime(df.index, format="mixed", errors="coerce")
|
|
df = df[~df.index.isna()]
|
|
df = df[~df.index.duplicated(keep="last")].sort_index()
|
|
return df
|
|
|
|
|
|
def frozen_syms() -> set:
|
|
"""Symbols with a frozen history copy in the overrides dir."""
|
|
if not FROZEN_DIR.is_dir():
|
|
return set()
|
|
return {f.stem[: -len("-history")] for f in FROZEN_DIR.glob("*-history.csv")}
|
|
|
|
|
|
def _panel_paths(root: Path, suffix: str) -> list[tuple[str, Path]]:
|
|
"""(symbol, path) for every panel file; frozen copies shadow data-root."""
|
|
frozen = frozen_syms()
|
|
out = []
|
|
for f in sorted(root.glob(f"*-{suffix}.csv")):
|
|
sym = f.name[: -len(f"-{suffix}.csv")]
|
|
if sym in frozen:
|
|
continue
|
|
out.append((sym, f))
|
|
for f in sorted(FROZEN_DIR.glob(f"*-{suffix}.csv")) if FROZEN_DIR.is_dir() else []:
|
|
out.append((f.name[: -len(f"-{suffix}.csv")], f))
|
|
return out
|
|
|
|
|
|
def _read_panel(root: Path, suffix: str, col: str) -> pd.DataFrame:
|
|
cols, syms = [], []
|
|
for sym, f in _panel_paths(root, suffix):
|
|
try:
|
|
df = _read_csv_clean(f)
|
|
except Exception:
|
|
continue
|
|
if col in df.columns:
|
|
cols.append(_apply_corrections(sym, suffix, df[col]))
|
|
syms.append(sym)
|
|
if not cols:
|
|
return pd.DataFrame()
|
|
out = pd.concat(cols, axis=1)
|
|
out.columns = syms
|
|
return out.sort_index()
|
|
|
|
|
|
def _read_name(f: Path) -> str:
|
|
try:
|
|
meta = json.loads(f.read_text())["chart"]["result"][0]["meta"]
|
|
return meta.get("longName", f.stem)
|
|
except Exception:
|
|
return f.stem
|
|
|
|
|
|
def _read_names(root: Path, max_files: int = 5000) -> dict:
|
|
files = sorted(root.glob("*.json"))[:max_files]
|
|
return {f.stem: _read_name(f) for f in files}
|
|
|
|
|
|
# ---------------------------------------------------------------- caching
|
|
# a bundle is ~2 GB resident; memoize the last one so Streamlit re-runs
|
|
# (every widget change!) don't reload the parquet panels
|
|
_BUNDLE_CACHE: dict[tuple, Bundle] = {}
|
|
_LOCK = threading.Lock()
|
|
|
|
# Background cache maintenance. The data root is written by an independent
|
|
# downloader (goget) while Streamlit serves requests; the two must not
|
|
# interfere. `load_bundle` therefore never refreshes on the request path:
|
|
# it serves the in-memory bundle immediately and, if the data dir changed,
|
|
# a background thread refreshes the cache at its own pace.
|
|
WATCH_INTERVAL = 5.0 # seconds between staleness scans
|
|
REFRESH_CADENCE = 30.0 # min seconds between actual cache rebuilds
|
|
_watchers: dict[tuple, threading.Thread] = {}
|
|
_gen: dict[tuple, int] = {} # (root, cache) -> bundle generation counter
|
|
def _scan_files(root: Path) -> dict:
|
|
"""{filename: (mtime_ns, size)} for every csv/json file in the data root
|
|
plus the frozen overrides (namespaced under 'frozen/')."""
|
|
out = {}
|
|
with os.scandir(root) as it:
|
|
for e in it:
|
|
if e.is_file() and e.name.endswith((".csv", ".json")):
|
|
st = e.stat()
|
|
out[e.name] = (st.st_mtime_ns, st.st_size)
|
|
if FROZEN_DIR.is_dir():
|
|
with os.scandir(FROZEN_DIR) as it:
|
|
for e in it:
|
|
if e.is_file() and e.name.endswith((".csv", ".json")):
|
|
st = e.stat()
|
|
out["frozen/" + e.name] = (st.st_mtime_ns, st.st_size)
|
|
if CORRECTIONS_DIR.is_dir():
|
|
with os.scandir(CORRECTIONS_DIR) as it:
|
|
for e in it:
|
|
if e.is_file() and e.name.endswith(".json"):
|
|
st = e.stat()
|
|
out["corrections/" + e.name] = (st.st_mtime_ns, st.st_size)
|
|
return out
|
|
|
|
|
|
def _manifest_path(cache: Path) -> Path:
|
|
return cache / "manifest.json"
|
|
|
|
|
|
def _cache_up_to_date(root: Path, cache: Path) -> bool:
|
|
"""Cheap staleness check: dir scan vs manifest.
|
|
|
|
A scan is a handful of milliseconds for ~16k files, so this is done on
|
|
every call (caching the scan result would mask data updates)."""
|
|
new = _scan_files(root)
|
|
try:
|
|
old = json.loads(_manifest_path(cache).read_text())
|
|
except Exception:
|
|
return False
|
|
return {f: tuple(v) for f, v in old.items()} == new
|
|
|
|
|
|
def _full_build(root: Path, cache: Path, cache_files: list) -> None:
|
|
"""(Re)read every CSV and write all parquet panels + manifest."""
|
|
adj = _read_panel(root, "history", "Adj Close")
|
|
close = _read_panel(root, "history", "Close")
|
|
div = _read_panel(root, "dividend", "Dividends").fillna(0.0)
|
|
capg = _read_panel(root, "capitalGain", "Capital Gains").fillna(0.0)
|
|
for df, f in zip((adj, close, div, capg), cache_files):
|
|
df.to_parquet(f)
|
|
_manifest_path(cache).write_text(json.dumps(_scan_files(root)))
|
|
|
|
|
|
def _apply_changes(root: Path, cache: Path, old: dict, new: dict) -> None:
|
|
"""Incremental refresh: re-read only files that were added/changed,
|
|
drop symbols whose files vanished, refresh names, rewrite the touched
|
|
parquet panels.
|
|
|
|
Per touched panel this is ONE read + ONE concat + ONE write (inserting
|
|
column-by-column into a ~500 MB frame is prohibitively slow).
|
|
"""
|
|
old = {f: tuple(v) for f, v in old.items()} # JSON round-trip turns tuples into lists
|
|
changed = [f for f, v in new.items() if old.get(f) != v]
|
|
removed = [f for f in old if f not in new]
|
|
if not changed and not removed:
|
|
return
|
|
|
|
def _suffix(fname: str):
|
|
for sfx in _SUFFIXES:
|
|
if fname.endswith(f"-{sfx}.csv"):
|
|
return sfx, fname[: -len(f"-{sfx}.csv")]
|
|
return None, None
|
|
|
|
def _path(fname: str) -> Path:
|
|
# frozen/ namespaced files live in the overrides dir
|
|
return FROZEN_DIR / fname[8:] if fname.startswith("frozen/") else root / fname
|
|
|
|
updates: dict[str, dict] = {k: {} for k in _PANELS} # panel -> {sym: series}
|
|
drops: set[str] = set()
|
|
# data-root files of frozen symbols are ignored: their canonical series
|
|
# lives in the overrides dir, which is tracked in the same scan
|
|
frozen = frozen_syms()
|
|
|
|
def _recompute_corrected(symfile: str) -> None:
|
|
"""corrections/{SYM}.json changed: re-derive the symbol's
|
|
dividend/capital-gain series from the base file + corrections."""
|
|
for sfx, colname, key in (("dividend", "Dividends", "div"),
|
|
("capitalGain", "Capital Gains", "capg")):
|
|
sym = None
|
|
p = None
|
|
for case in (symfile[:-5].lower(), symfile[:-5].upper()):
|
|
for base in (FROZEN_DIR, root):
|
|
q = base / f"{case}-{sfx}.csv"
|
|
if q.exists():
|
|
sym, p = case, q
|
|
break
|
|
if sym:
|
|
break
|
|
if p is None:
|
|
continue
|
|
try:
|
|
df = _read_csv_clean(p)
|
|
except Exception:
|
|
continue
|
|
if colname not in df.columns:
|
|
continue
|
|
updates[key][sym] = _apply_corrections(sym, sfx, df[colname]) \
|
|
.fillna(0.0)
|
|
|
|
for fname in changed:
|
|
if fname.startswith("corrections/"):
|
|
_recompute_corrected(fname[12:])
|
|
continue
|
|
sfx, sym = _suffix(fname)
|
|
if sfx is None:
|
|
continue
|
|
if not fname.startswith("frozen/") and sym in frozen:
|
|
continue
|
|
try:
|
|
df = _read_csv_clean(_path(fname))
|
|
except Exception:
|
|
continue # e.g. a file mid-download
|
|
if sfx == "history":
|
|
if "Adj Close" in df.columns:
|
|
updates["adj"][sym] = df["Adj Close"]
|
|
if "Close" in df.columns:
|
|
updates["close"][sym] = df["Close"]
|
|
elif sfx == "dividend" and "Dividends" in df.columns:
|
|
updates["div"][sym] = _apply_corrections(sym, sfx, df["Dividends"]) \
|
|
.fillna(0.0)
|
|
elif sfx == "capitalGain" and "Capital Gains" in df.columns:
|
|
updates["capg"][sym] = _apply_corrections(sym, sfx, df["Capital Gains"]) \
|
|
.fillna(0.0)
|
|
for fname in removed:
|
|
sfx, sym = _suffix(fname)
|
|
if fname.startswith("corrections/"):
|
|
_recompute_corrected(fname[12:])
|
|
continue
|
|
if fname.startswith("frozen/"):
|
|
continue # frozen files are not data-root symbols
|
|
if sym in frozen:
|
|
continue # data-root file of a frozen symbol vanished; the
|
|
# frozen copy remains authoritative
|
|
sfx, sym = _suffix(fname)
|
|
if sfx is not None:
|
|
drops.add(sym)
|
|
|
|
# names: only the json metadata files that changed
|
|
j_changed = [f for f in changed if f.endswith(".json")]
|
|
j_removed = [f for f in removed if f.endswith(".json")]
|
|
if j_changed or j_removed:
|
|
names_path = cache / "names.json"
|
|
try:
|
|
names = json.loads(names_path.read_text())
|
|
except Exception:
|
|
names = {}
|
|
for f in j_removed:
|
|
names.pop(f[:-5], None)
|
|
for f in j_changed:
|
|
names[f[:-5]] = _read_name(root / f)
|
|
names_path.write_text(json.dumps(names))
|
|
|
|
for k in _PANELS:
|
|
upd = {s: ser for s, ser in updates[k].items() if s not in drops}
|
|
if not upd and not drops:
|
|
continue
|
|
panel = pd.read_parquet(cache / f"panel_{k}.parquet")
|
|
old = {s: panel[s] for s in upd if s in panel.columns}
|
|
drop_cols = [s for s in (drops | set(upd)) if s in panel.columns]
|
|
if drop_cols:
|
|
panel = panel.drop(columns=drop_cols)
|
|
if upd:
|
|
extras = []
|
|
for sym, s in upd.items():
|
|
s = s.rename(sym)
|
|
if sym in old:
|
|
# new values win where present; keep old values for dates
|
|
# the new file doesn't cover (e.g. a partial re-download)
|
|
uidx = panel.index.union(s.index)
|
|
s = s.reindex(uidx).combine_first(old[sym].reindex(uidx))
|
|
extras.append(s)
|
|
extra = pd.concat(extras, axis=1)
|
|
extra.columns = list(upd)
|
|
panel = pd.concat([panel, extra], axis=1).sort_index()
|
|
panel.to_parquet(cache / f"panel_{k}.parquet")
|
|
|
|
|
|
def up_to_date(root: Path = DEFAULT_ROOT, cache: Path = CACHE_DIR) -> bool:
|
|
"""Cheap check whether the cache matches the data dir right now."""
|
|
return _cache_up_to_date(Path(root), Path(cache))
|
|
|
|
|
|
def generation(root: Path = DEFAULT_ROOT, cache: Path = CACHE_DIR) -> int:
|
|
"""Monotonic counter, bumped each time the bundle is replaced.
|
|
|
|
Feed it into st.cache_data keys so cached computations invalidate when
|
|
the data changes."""
|
|
return _gen.get((str(root), str(cache)), 0)
|
|
|
|
|
|
def refresh(root: Path = DEFAULT_ROOT, cache: Path = CACHE_DIR,
|
|
rebuild: bool = False) -> Bundle:
|
|
"""Synchronously refresh the parquet cache if stale and return the bundle.
|
|
|
|
The cache tracks the data dir via a manifest of (mtime, size) per file:
|
|
changed/added files are re-read and merged in incrementally, removed
|
|
files are dropped, so an updated download refreshes in seconds instead
|
|
of the ~1 min full rebuild. `rebuild=True` forces the full rebuild.
|
|
"""
|
|
root, cache = Path(root), Path(cache)
|
|
key = (str(root), rebuild)
|
|
with _LOCK:
|
|
hit = _BUNDLE_CACHE.get(key)
|
|
if hit is not None and not rebuild and _cache_up_to_date(root, cache):
|
|
return hit
|
|
cache.mkdir(parents=True, exist_ok=True)
|
|
cache_files = [cache / f"panel_{k}.parquet" for k in _PANELS]
|
|
cache_ok = all(f.exists() for f in cache_files)
|
|
|
|
if rebuild or not cache_ok:
|
|
_full_build(root, cache, cache_files)
|
|
else:
|
|
try:
|
|
old = json.loads(_manifest_path(cache).read_text())
|
|
except Exception:
|
|
old = None
|
|
new = _scan_files(root)
|
|
if old is None:
|
|
# parquet without a manifest: state unknown -> rebuild once
|
|
_full_build(root, cache, cache_files)
|
|
else:
|
|
_apply_changes(root, cache, old, new)
|
|
_manifest_path(cache).write_text(json.dumps(new))
|
|
|
|
panels = {k: pd.read_parquet(f) for k, f in zip(_PANELS, cache_files)}
|
|
names_path = cache / "names.json"
|
|
if names_path.exists():
|
|
names = json.loads(names_path.read_text())
|
|
else:
|
|
names = _read_names(root)
|
|
names_path.write_text(json.dumps(names))
|
|
bundle = Bundle(panels["adj"], panels["close"], panels["div"],
|
|
panels["capg"], names)
|
|
# keep only the newest bundle: there is no room for two
|
|
_BUNDLE_CACHE.clear()
|
|
_BUNDLE_CACHE[key] = bundle
|
|
_gen[(str(root), str(cache))] = _gen.get((str(root), str(cache)), 0) + 1
|
|
return bundle
|
|
|
|
|
|
def load_bundle(root: Path = DEFAULT_ROOT, cache: Path = CACHE_DIR,
|
|
rebuild: bool = False) -> Bundle:
|
|
"""Return the current bundle without blocking on data changes.
|
|
|
|
The first call (and `rebuild=True`) build synchronously — there is
|
|
nothing to serve yet. Otherwise the in-memory bundle is returned
|
|
immediately; if the data dir has changed, a background thread refreshes
|
|
the cache on its own schedule (every REFRESH_CADENCE seconds while the
|
|
downloader keeps writing), so serving requests and ingesting new data
|
|
run independently.
|
|
"""
|
|
root, cache = Path(root), Path(cache)
|
|
key = (str(root), rebuild)
|
|
with _LOCK:
|
|
hit = _BUNDLE_CACHE.get(key)
|
|
if hit is None:
|
|
return refresh(root, cache, rebuild)
|
|
if rebuild:
|
|
return refresh(root, cache, True)
|
|
if not _cache_up_to_date(root, cache):
|
|
_ensure_watcher(root, cache, key)
|
|
return hit
|
|
|
|
|
|
def _ensure_watcher(root: Path, cache: Path, key: tuple) -> None:
|
|
wk = (str(root), str(cache))
|
|
t = _watchers.get(wk)
|
|
if t is not None and t.is_alive():
|
|
return
|
|
t = threading.Thread(target=_watcher_loop, args=(root, cache, key),
|
|
daemon=True)
|
|
_watchers[wk] = t
|
|
t.start()
|
|
|
|
|
|
def _watcher_loop(root: Path, cache: Path, key: tuple) -> None:
|
|
last = 0.0
|
|
while True:
|
|
time.sleep(WATCH_INTERVAL)
|
|
with _LOCK:
|
|
active = _BUNDLE_CACHE.get(key) is not None
|
|
if not active:
|
|
return # bundle evicted (another key loaded) -> stop
|
|
if _cache_up_to_date(root, cache):
|
|
continue
|
|
now = time.monotonic()
|
|
if now - last < REFRESH_CADENCE:
|
|
continue
|
|
try:
|
|
refresh(root, cache)
|
|
except Exception:
|
|
continue # transient (e.g. partial read); retry next tick
|
|
last = time.monotonic()
|
|
|
|
|
|
def search_symbols(bundle: Bundle, query: str, limit: int = 200) -> list[str]:
|
|
"""Case-insensitive search over symbols and names.
|
|
|
|
Ranked: exact symbol, symbol prefix, name prefix, symbol substring,
|
|
name substring.
|
|
"""
|
|
q = query.lower()
|
|
buckets: dict[int, list[str]] = {}
|
|
for s in bundle.adj.columns:
|
|
sl, nl = s.lower(), bundle.names.get(s, "").lower()
|
|
if sl == q:
|
|
buckets.setdefault(0, []).append(s)
|
|
elif sl.startswith(q):
|
|
buckets.setdefault(1, []).append(s)
|
|
elif nl.startswith(q):
|
|
buckets.setdefault(2, []).append(s)
|
|
elif q in sl:
|
|
buckets.setdefault(3, []).append(s)
|
|
elif q in nl:
|
|
buckets.setdefault(4, []).append(s)
|
|
out: list[str] = []
|
|
for k in sorted(buckets):
|
|
out.extend(buckets[k])
|
|
return out[:limit]
|