data: manifest-based incremental refresh of the parquet cache
The cache now tracks every file in the data dir (mtime_ns + size) in
.cache/manifest.json. On load, a directory scan is compared against the
manifest:
- changed/added files are re-read and merged into the parquet panels
(one read + one concat + one write per touched panel; new values
win where present, old values kept where the new file is short)
- removed files drop their symbols (and names)
- an up-to-date cache is a ~30 ms memo hit
Measured on the real 4k-symbol set: full build 54 s, refresh of
5 modified + 1 added + 1 removed files 3.4 s. No scan TTL (a scan is
a few ms); a previous 5 s scan cache masked data updates.
Tests: tests/test_data.py (11 checks) added as step 1 of run_tests.sh.
This commit is contained in:
parent
d8703a7a63
commit
cba7291676
|
|
@ -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). |
|
||||
|
|
|
|||
248
data.py
248
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]:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
120
tests/test_data.py
Normal file
120
tests/test_data.py
Normal file
|
|
@ -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())
|
||||
Loading…
Reference in New Issue
Block a user