Greg Pomerantz
02aa750a06
230 tickers whose Yahoo chart responses now come back without a timestamp array (terminated/merged funds): goget overwrites the .json on every pass while ohlc.Conv skips the write, leaving the old CSVs as the last known series. Snapshot them into overrides/frozen/ (git-tracked, audited in reports/stale-funds.md) and make data.py prefer the frozen copies and ignore any future data-root rewrite/delete for those symbols, so the final series survives future goget runs. The cache manifest now covers the overrides dir too; incremental refresh skips data-root files of frozen symbols.
411 lines
16 KiB
Python
411 lines
16 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"
|
|
|
|
_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(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)
|
|
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()
|
|
for fname in changed:
|
|
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] = df["Dividends"].fillna(0.0)
|
|
elif sfx == "capitalGain" and "Capital Gains" in df.columns:
|
|
updates["capg"][sym] = df["Capital Gains"].fillna(0.0)
|
|
for fname in removed:
|
|
sfx, sym = _suffix(fname)
|
|
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]
|