gmp/f
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
d5a7bb70b0
f/data.py
Greg Pomerantz 5e5a10d725 Fix Yahoo total-distribution double count (1,957 syms) 
Yahoo's dividend endpoint returns the fund's TOTAL per-share distribution
(dividend + capital gain) for many share classes while the capitalGains
endpoint returns the cap-gain portion separately; summing both double-counted
cap gains and broke pre/post-tax comparability (after-tax engine could beat
the adj-based pre-tax return: impossible).

Add _fix_total_distributions to the bundle-assembly invariants
(layout-agnostic, survives re-downloads): on same-date div>capg events, the
exact Yahoo-adj implied distribution (d = P - Q*A_{t-1}/A_t) must match the
div-file amount (not the sum) on >=3 dates before rewriting div -= capg.
Self-validating: genuine separate same-date distributions are untouched.
Rewrote 17,900 cells on 1,957 symbols.

Pool check (scripts/check_adj_consistency.py): all 14 symbols now
TR(close+events) == TR(adj) to 0.01pt and post < pre with plausible drags.

Tests: rewrite case, no-rewrite case, idempotency (tests/test_data.py).
2026-09-01 16:23:03 -04:00

688 lines
28 KiB
Python
Raw Blame History

"""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 numpy as np
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"
# Backup of the dividend/capitalGain event files, used as a fallback when a
# re-download no longer carries those events. In 2026 Yahoo changed its event
# feed: for some funds it stopped returning capitalGain events entirely and
# for terminated funds it returns no events at all, so a plain re-download
# overwrites good historical event CSVs with empty ones. The backup holds
# the last known-good content; see scripts/backup_events.py.
EVENT_BACKUP_DIR = Path(__file__).parent / "overrides" / "event-backup"
# 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...},
# "dedup": [["2022-12-13", 4.841], ...],
# "drop_capg_copy": ["2023-12-21", ...],
# "as_of": "2026-08-31", "source": "<filing URL>", "note": "..."}
#
# The two cross-file ops express the Yahoo double-listing invariant and are
# LAYOUT-AGNOSTIC (applied at bundle assembly, idempotent): Yahoo has changed
# which file carries a distribution between downloads (e.g. in 2026 it stopped
# returning capitalGain events entirely for some funds while the dividend
# stream kept the year-end row), so file-specific removes would silently break
# on the next re-download.
# dedup: [[date, amount]] the (date, amount) row may appear in
# both files; keep at most one (the capitalGain copy if both present).
# drop_capg_copy: [date] the capitalGain row on that date is a
# spurious copy of the dividend row; drop it (any amount).
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 dedupe_event_rows(corr: dict, div: list[tuple[str, float]],
capg: list[tuple[str, float]]
) -> tuple[list[tuple[str, float]], list[tuple[str, float]]]:
"""Pure list-level version of apply_invariants (no file access), shared
by verify_official and unit tests."""
dups = corr.get("dedup") or []
drops = corr.get("drop_capg_copy") or []
if not dups and not drops:
return div, capg
ck = {(d, round(v, 9)) for d, v in capg}
# dedup: drop the dividend copy only when the capitalGain copy is present
out_div = [r for r in div
if not any(r[0] == d and round(r[1], 9) == round(a, 9)
and (d, round(a, 9)) in ck
for d, a in dups)]
# drop_capg_copy: the capg row on that date is the spurious one
out_capg = [r for r in capg if r[0] not in drops]
return out_div, out_capg
def _fix_total_distributions(div: pd.DataFrame, capg: pd.DataFrame,
close: pd.DataFrame, adj: pd.DataFrame) -> list[str]:
"""Fix Yahoo's "dividend file carries the TOTAL distribution" quirk.
For many fund share classes Yahoo's dividend endpoint reports the fund's
TOTAL per-share distribution (dividend + capital gain) while the
capitalGains endpoint reports the cap-gain portion separately; summing
both files double-counts the cap-gain part and breaks pre/post-tax
comparability (the after-tax engine can even beat the adj-based pre-tax
return, which is impossible for consistent data).
Detection is layout-agnostic and self-validating: on dates where both
files have an event and the dividend amount strictly exceeds the cap-gain
amount, Yahoo's adj factor steps by an implicit total distribution (the
close/adj ratio only steps on distribution dates). If that implicit
amount matches the dividend-file amount (not the sum) on most overlap
dates, the dividend file is the total and we rewrite div -= capg there,
leaving div+capg == total. Symbols with genuine separate same-date
distributions (adj matches the sum) are left untouched.
Returns the list of rewritten symbols.
"""
fixed: list[str] = []
for sym in div.columns:
if sym not in capg.columns or sym not in close.columns or sym not in adj.columns:
continue
dv = div[sym]
cg = capg[sym].reindex(div.index).fillna(0.0)
mask = (dv > 0) & (cg > 0) & (dv > cg * 1.001)
n = int(mask.sum())
if n < 3:
continue
c, a = close[sym], adj[sym]
if (a.fillna(0) == 0).all() or (c.fillna(0) == 0).all():
continue
# Yahoo adj convention: A_t = A_{t-1} * Q_t / (P_{t-1} - d) =>
# exact implied total distribution: d = P - Q * (A_{t-1} / A_t)
imp = (c.shift(1) - c * (a.shift(1) / a)).reindex(div.index)
d_s = dv[mask]
g_s = cg[mask]
imp = imp[mask]
ref = np.maximum(imp.abs(), d_s.abs())
ref = ref.replace(0, np.nan)
m_tot = (imp - d_s).abs() <= 0.05 * ref
m_both = (imp - (d_s + g_s)).abs() <= 0.05 * ref
if int(m_tot.sum()) >= 3 and int(m_tot.sum()) > int(m_both.sum()):
div.loc[mask, sym] = d_s - g_s
fixed.append(sym)
return fixed
def apply_invariants(div: pd.DataFrame, capg: pd.DataFrame,
adj: pd.DataFrame | None = None,
close: pd.DataFrame | None = None) -> None:
"""Cross-file correction invariants (Yahoo double-listing), applied at
bundle assembly so they hold for full and incremental builds and survive
re-downloads that move a row between the dividend/capitalGain files.
Idempotent: panels keep the pre-invariant state, this re-derives it."""
for f in sorted(CORRECTIONS_DIR.glob("*.json")):
sym = f.stem.lower()
try:
corr = json.loads(f.read_text())
except Exception:
continue
dups = corr.get("dedup") or []
drops = corr.get("drop_capg_copy") or []
if not dups and not drops:
continue
if sym in capg.columns:
for d in drops:
ts = pd.Timestamp(d)
if ts in capg.index and abs(capg.at[ts, sym]) > 1e-9:
capg.at[ts, sym] = 0.0
if sym in div.columns and sym in capg.columns:
for d, a in dups:
ts = pd.Timestamp(d)
if ts in div.index and ts in capg.index and \
abs(div.at[ts, sym] - a) < 1e-9 and \
abs(capg.at[ts, sym] - a) < 1e-9:
div.at[ts, sym] = 0.0
if adj is not None and close is not None:
_fix_total_distributions(div, capg, close, adj)
def _apply_corrections(sym: str, suffix: str, ser: pd.Series,
col: str | None = None) -> pd.Series:
if suffix == "history":
# {date: {col: value}} — patch individual OHLC cells (e.g. a
# misaligned Yahoo adj-close spike). Only touches the requested col.
ops = _load_corrections(sym).get("history")
if not ops or col is None:
return ser
d = {ts: float(v) for ts, v in ser.items()}
for dstr, cells in ops.items():
ts = pd.Timestamp(dstr)
if ts in d and col in cells:
d[ts] = float(cells[col])
return pd.Series(d).sort_index()
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 _has_data_rows(p: Path) -> bool:
"""True if the csv has at least one row beyond the header."""
try:
with open(p, errors="replace") as fh:
if not fh.readline():
return False
return fh.readline() != ""
except OSError:
return False
def event_file(sym: str, sfx: str, root: Path) -> Path:
"""Effective event file (dividend/capitalGain): frozen copy, else the
data-root file while it still carries rows, else the backup, else the
(possibly emptied by a re-download) data-root file."""
fr = FROZEN_DIR / f"{sym}-{sfx}.csv"
if fr.exists():
return fr
rt = root / f"{sym}-{sfx}.csv"
if _has_data_rows(rt):
return rt
bk = EVENT_BACKUP_DIR / f"{sym}-{sfx}.csv"
if bk.exists() and _has_data_rows(bk):
return bk
return rt
def _panel_paths(root: Path, suffix: str) -> list[tuple[str, Path]]:
"""(symbol, path) for every panel file; frozen copies shadow data-root,
and for event files the backup dir is the fallback (see event_file)."""
frozen = frozen_syms()
if suffix in ("dividend", "capitalGain"):
syms = sorted({f.name[: -len(f"-{suffix}.csv")]
for f in root.glob(f"*-{suffix}.csv")}
| ({f.name[: -len(f"-{suffix}.csv")]
for f in EVENT_BACKUP_DIR.glob(f"*-{suffix}.csv")}
if EVENT_BACKUP_DIR.is_dir() else set())
| ({f.name[: -len(f"-{suffix}.csv")]
for f in FROZEN_DIR.glob(f"*-{suffix}.csv")}
if FROZEN_DIR.is_dir() else set()))
return [(sym, event_file(sym, suffix, root))
for sym in syms if sym not in frozen
or (FROZEN_DIR / f"{sym}-{suffix}.csv").exists()]
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], 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()
authoritative: set[str] = set() # corrections-derived series: replace,
# don't combine (deletions must stick)
# 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 (("history", "Adj Close", "adj"),
("history", "Close", "close"),
("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
ser = _apply_corrections(sym, sfx, df[colname], colname)
updates[key][sym] = ser if sfx == "history" else ser.fillna(0.0)
authoritative.add(sym)
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] = _apply_corrections(
sym, "history", df["Adj Close"], "Adj Close")
if "Close" in df.columns:
updates["close"][sym] = _apply_corrections(
sym, "history", df["Close"], "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:
uidx = panel.index.union(s.index)
if sym in authoritative:
# corrections-derived: a date missing from the new
# series was intentionally removed, not "not covered"
s = s.reindex(uidx).fillna(0.0)
else:
# new values win where present; keep old values for
# dates the new file doesn't cover (partial download)
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))
apply_invariants(panels["div"], panels["capg"],
panels["adj"], panels["close"])
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]
Reference in New Issue View Git Blame Copy Permalink