diff --git a/README.md b/README.md index 2639956..e3dcde4 100644 --- a/README.md +++ b/README.md @@ -12,13 +12,16 @@ python3 -m venv .venv ``` First run builds a parquet cache in `.cache/` (~1 min for 4k symbols); -later runs load in well under a second. +later runs load in well under a second. The cache tracks the data dir +per-file (mtime + size in `.cache/manifest.json`), so when the download +is updated, only the changed/added/removed symbols are re-read — a +partial refresh takes seconds instead of a full ~1 min rebuild. ## Modules | Module | Purpose | |-----------------|---------| -| `data.py` | Ingest `{sym}-history/dividend/capitalGain.csv` -> cached parquet panels (date x symbol). `Adj Close` already includes distributions, so it drives pre-tax total returns. | +| `data.py` | Ingest `{sym}-history/dividend/capitalGain.csv` -> cached parquet panels (date x symbol), with manifest-based incremental refresh when the data dir changes. `Adj Close` already includes distributions, so it drives pre-tax total returns. | | `metrics.py` | Total/annualized return, vol, Sharpe, Sortino, max drawdown, Calmar, CAPM beta/alpha. Pure pandas, all transparent. | | `portfolio.py` | Weighted portfolios with drift and periodic rebalancing to target weights (`1W/1ME/QE/YE`), one-way cost in bps. Spec grammar: commas join the elements of ONE portfolio (`SYM` or `SYM:w`, bare = equal weight), spaces separate DISTINCT symbols/portfolios (`parse_items`). | | `tax.py` | Simplified DAS after-tax engine: FIFO lots, 365-day long/short split, separate LT/ST/dividend rates. Headline curve = what you keep if you **sell everything today** (unrealized gains taxed daily by lot age). | diff --git a/data.py b/data.py index b5f142f..7d07f0b 100644 --- a/data.py +++ b/data.py @@ -18,6 +18,8 @@ First call builds a parquet cache (default: .cache/) in ~1 min for from __future__ import annotations import json +import os +import threading from dataclasses import dataclass from pathlib import Path @@ -26,6 +28,9 @@ import pandas as pd DEFAULT_ROOT = Path("~/prog/fin/stocks").expanduser() CACHE_DIR = Path(__file__).parent / ".cache" +_SUFFIXES = ("history", "dividend", "capitalGain") +_PANELS = ("adj", "close", "div", "capg") + @dataclass class Bundle: @@ -36,19 +41,24 @@ class Bundle: 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 _read_panel(root: Path, suffix: str, col: str) -> pd.DataFrame: - cols = [] - syms = [] + cols, syms = [], [] for f in sorted(root.glob(f"*-{suffix}.csv")): - sym = f.name[: -len(f"-{suffix}.csv")] - 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() - if col not in df.columns: + try: + df = _read_csv_clean(f) + except Exception: continue - cols.append(df[col]) - syms.append(sym) + if col in df.columns: + cols.append(df[col]) + syms.append(f.name[: -len(f"-{suffix}.csv")]) if not cols: return pd.DataFrame() out = pd.concat(cols, axis=1) @@ -56,58 +66,194 @@ def _read_panel(root: Path, suffix: str, col: str) -> pd.DataFrame: 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: - names = {} - for f in root.glob("*.json"): - try: - meta = json.loads(f.read_text())["chart"]["result"][0]["meta"] - names[f.stem] = meta.get("longName", f.stem) - except Exception: - names[f.stem] = f.stem - return names + 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() +def _scan_files(root: Path) -> dict: + """{filename: (mtime_ns, size)} for every csv/json file in the data root.""" + 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) + 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 + + updates: dict[str, dict] = {k: {} for k in _PANELS} # panel -> {sym: series} + drops: set[str] = set() + for fname in changed: + sfx, sym = _suffix(fname) + if sfx is None: + continue + try: + df = _read_csv_clean(root / 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 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 load_bundle(root: Path = DEFAULT_ROOT, cache: Path = CACHE_DIR, rebuild: bool = False) -> Bundle: - """Load the full dataset, using (or building) a parquet cache.""" - key = (str(root), rebuild) - hit = _BUNDLE_CACHE.get(key) - if hit is not None: - return hit - cache = Path(cache) - cache.mkdir(parents=True, exist_ok=True) - cache_files = [cache / f"panel_{k}.parquet" for k in ("adj", "close", "div", "capg")] + """Load the full dataset, using (or building) a parquet cache. - if not rebuild and all(f.exists() for f in cache_files): - adj = pd.read_parquet(cache_files[0]) - close = pd.read_parquet(cache_files[1]) - div = pd.read_parquet(cache_files[2]) - capg = pd.read_parquet(cache_files[3]) - else: - 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) - adj.to_parquet(cache_files[0]) - close.to_parquet(cache_files[1]) - div.to_parquet(cache_files[2]) - capg.to_parquet(cache_files[3]) - 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(adj, close, div, capg, names) - # keep only the newest bundle: there is no room for two - _BUNDLE_CACHE.clear() - _BUNDLE_CACHE[key] = bundle - return 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 + return bundle def search_symbols(bundle: Bundle, query: str, limit: int = 200) -> list[str]: diff --git a/run_tests.sh b/run_tests.sh index 6df6dda..34a66c2 100755 --- a/run_tests.sh +++ b/run_tests.sh @@ -9,11 +9,15 @@ cd "$(dirname "$0")" fail=0 -echo "=== 1/2 app tests (AppTest, no browser) ===" +echo "=== 1/3 data cache tests (incremental refresh) ===" +.venv/bin/python tests/test_data.py || fail=1 +echo + +echo "=== 2/3 app tests (AppTest, no browser) ===" .venv/bin/python tests/test_app.py || fail=1 echo -echo "=== 2/2 browser e2e (Playwright, needs the server) ===" +echo "=== 3/3 browser e2e (Playwright, needs the server) ===" if curl -s -m 3 -o /dev/null http://localhost:8599/healthz; then .venv/bin/python tests/test_e2e_browser.py || fail=1 else diff --git a/tests/test_data.py b/tests/test_data.py new file mode 100644 index 0000000..ca85f19 --- /dev/null +++ b/tests/test_data.py @@ -0,0 +1,120 @@ +"""Data-cache tests: incremental refresh from changed data files. + +Uses a small synthetic data root in a temp dir — no real data, no server. +Run: .venv/bin/python tests/test_data.py +""" +from __future__ import annotations + +import json +import os +import pathlib +import shutil +import sys +import tempfile + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +import data as D # noqa: E402 + +PASS, FAIL = 0, 0 + + +def check(name: str, cond: bool, extra: str = "") -> None: + global PASS, FAIL + if cond: + PASS += 1 + print(f" ok {name}", flush=True) + else: + FAIL += 1 + print(f" FAIL {name} {extra}", flush=True) + + +def write_history(root: pathlib.Path, sym: str, last_price: float | None = None, + days: int = 10) -> pathlib.Path: + """Prices 100.00, 101.00, ..., 109.00 (last day overridable).""" + f = root / f"{sym}-history.csv" + lines = ["Date,Open,High,Low,Close,Adj Close,Volume"] + for i in range(days): + p = last_price if (last_price is not None and i == days - 1) \ + else 100.0 + i + lines.append(f"2024-01-{i + 1:02d},{p:.2f},{p:.2f},{p:.2f},{p:.2f},{p:.2f},1000") + f.write_text("\n".join(lines) + "\n") + return f + + +def touch(p: pathlib.Path) -> None: + """Bump mtime so the manifest sees the file as changed.""" + st = p.stat() + os.utime(p, ns=(st.st_atime_ns + 10**9, st.st_mtime_ns + 10**9)) + + +def reload(d: pathlib.Path) -> D.Bundle: + return D.load_bundle(root=d["root"], cache=d["cache"]) + + +def main() -> int: + tmp = pathlib.Path(tempfile.mkdtemp(prefix="fdatatest-")) + root, cache = tmp / "stocks", tmp / "cache" + root.mkdir() + d = {"root": root, "cache": cache} + try: + write_history(root, "aaa") + write_history(root, "bbb") + write_history(root, "ccc") + (root / "aaa-dividend.csv").write_text("Date,Dividends\n2024-01-05,0.5\n") + (root / "bbb.json").write_text( + json.dumps({"chart": {"result": [{"meta": {"longName": "Bee Bee Corp"}}]}})) + + b = reload(d) + check("initial full build", list(b.adj.columns) == ["aaa", "bbb", "ccc"] + and b.adj["aaa"].iloc[-1] == 109.0, str(list(b.adj.columns))) + check("names from json", b.names.get("bbb") == "Bee Bee Corp") + check("dividend panel", b.div["aaa"].loc["2024-01-05"] == 0.5) + + # --- modified file: last price changes + f = write_history(root, "aaa", last_price=200.0) + touch(f) + b = reload(d) + check("changed file re-read incrementally", b.adj["aaa"].iloc[-1] == 200.0) + check("untouched symbols intact", b.adj["bbb"].iloc[-1] == 109.0) + + # --- new file: symbol added + write_history(root, "ddd") + b = reload(d) + check("new symbol added", "ddd" in b.adj.columns) + + # --- deleted file: symbol dropped + (root / "ccc-history.csv").unlink() + b = reload(d) + check("deleted symbol dropped", "ccc" not in b.adj.columns + and "ccc" not in b.div.columns) + + # --- name file added for a new symbol + (root / "ddd.json").write_text( + json.dumps({"chart": {"result": [{"meta": {"longName": "Dee Dee"}}]}})) + b = reload(d) + check("name added incrementally", b.names.get("ddd") == "Dee Dee") + + # --- dividend file modified + (root / "aaa-dividend.csv").write_text( + "Date,Dividends\n2024-01-05,0.5\n2024-01-06,1.25\n") + b = reload(d) + check("dividend change picked up", b.div["aaa"].loc["2024-01-06"] == 1.25) + + # --- no-change reload: cheap, no rewrite + before = (cache / "panel_adj.parquet").stat().st_mtime_ns + b = reload(d) + after = (cache / "panel_adj.parquet").stat().st_mtime_ns + check("no-op reload does not rewrite parquet", before == after) + + # --- forced full rebuild still works + b = D.load_bundle(root=root, cache=cache, rebuild=True) + check("forced full rebuild", list(b.adj.columns) == ["aaa", "bbb", "ddd"]) + finally: + shutil.rmtree(tmp, ignore_errors=True) + print(f"\n{PASS} passed, {FAIL} failed") + return 1 if FAIL else 0 + + +if __name__ == "__main__": + sys.exit(main())