Stock & Portfolio Analyzer: full UI rework

- single spec grammar for symbol and benchmark fields: commas join one
  portfolio (MSFT:0.6,V:0.4), spaces separate distinct symbols/portfolios;
  both fields accept one or many entries
- benchmarks simulated with the same scheme/cost/tax rules; per-benchmark
  beta/alpha columns; after-tax benchmark curves
- global Curve mode (pre/after/both) above the tabs; clean names in
  single-curve mode
- live updates: field commits on Enter/blur, page recomputes per rerun;
  portfolio+tax sims cached (st.cache_data); plotly.js from CDN (4.6MB ->
  browser-cached) with F_INLINE_PLOTLY=1 offline fallback
- chart: legend underneath, solid lines, pan sticks to data edges
  (width-preserving), zoom edge-clamped
- inputs persist in settings.json across reloads/restarts/devices
- tests: tests/test_app.py (AppTest) + tests/test_e2e_browser.py
  (Playwright) via ./run_tests.sh
This commit is contained in:
Greg Pomerantz 2026-08-24 16:05:27 -04:00
commit d8703a7a63
16 changed files with 1969 additions and 0 deletions

5
.gitignore vendored Normal file
View File

@ -0,0 +1,5 @@
.venv/
__pycache__/
*.pyc
.cache/
settings.json

64
README.md Normal file
View File

@ -0,0 +1,64 @@
# Stock & Portfolio Analyzer
Interactive tool for analyzing individual securities and portfolios
against local Yahoo Finance dumps (`~/prog/fin/stocks`, ~4k symbols).
## Quick start
```bash
python3 -m venv .venv
.venv/bin/pip install -r requirements.txt
./run.sh # serves the UI on the fixed port 8599 (http://localhost:8599)
```
First run builds a parquet cache in `.cache/` (~1 min for 4k symbols);
later runs load in well under a second.
## 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. |
| `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). |
| `chart_widget.py` | Self-contained plotly.js chart in an iframe: mouse zoom/pan, x clamped to the data, view edges snapped to first/last data points with day-precise labels, y tight-fit, every line re-based to 1.0 at the left edge. |
| `portfolios.py` | Saved portfolio definitions in `portfolios.json` (name, spec, scheme, cost). |
| `settings.json` | Persisted UI inputs (symbol/benchmark specs, scheme, costs, tax rates, period, curve/window mode) — restored on every page load and server restart; delete to reset. |
| `app.py` | Streamlit UI: single "symbol or portfolio" spec field (page updates as soon as the input is valid; unknown symbols get click-to-fix "did you mean" suggestions) + a benchmark box with the same grammar (one benchmark per line; a line is a single symbol or a comma-joined portfolio, simulated with the same scheme/cost/tax rules — pre- and after-tax curves, first one drives beta/alpha), scheme/costs/tax rates, save + load/compare/delete portfolios (overlaid pre/after-tax curves), curve toggle (both / pre-tax only / after-tax only), stats table, allocation, per-year tax detail. |
## Development
- **Run**: `./run.sh` → http://localhost:8599 (fixed port; no-ops if a
server is already running). The chart loads plotly.js from a CDN; for
fully offline use set `F_INLINE_PLOTLY=1` in `run.sh`.
- **Test**: `./run_tests.sh`
1. `tests/test_app.py` — app-level tests via Streamlit AppTest (no
browser). Memory: one data bundle is ~2.3 GB, so this process keeps
at most ONE AppTest alive (see its header comment).
2. `tests/test_e2e_browser.py` — Playwright + headless Chromium driving
the real page with real keystrokes; needs the server running on 8599.
One-time setup: `.venv/bin/pip install playwright` and
`.venv/bin/python -m playwright install chromium`.
- **Gotchas**
- Streamlit caches imported modules per process: **restart the server**
after editing any `.py` (kill the old one first — `run.sh` refuses to
double-start).
- `st.cache_data` caches the portfolio + tax simulations: they recompute
only when symbols/scheme/cost/tax rates change, not on window or curve
toggles.
- `settings.json` (gitignored) persists UI inputs across reloads and
restarts; delete it to reset. Saved portfolios live in `portfolios.json`.
- Data cache: `.cache/*.parquet`; rebuild via the sidebar checkbox
(first build ~1 min for ~4k symbols).
## Known simplifications (roadmap)
- No loss carryover or carryforward across years; no wash-sale rules.
- Distributed capital gains taxed entirely at the long-term rate.
- Single (federal-like) tax bracket; no state taxes, no AMT.
- Equal treatment of benchmark for beta/alpha (CAPM, rf = 0 by default).
Ideas: vectorbt sweeps over rebalance schemes, NiceGUI/Textual frontend,
empyrical-reloaded metrics, monthly (not yearly) loss netting, tax-loss
harvesting simulation.

434
app.py Normal file
View File

@ -0,0 +1,434 @@
"""Streamlit UI for stock/portfolio analysis.
Run: streamlit run app.py
"""
from __future__ import annotations
import json
import re
from pathlib import Path
import pandas as pd
import plotly.graph_objects as go
import streamlit as st
import metrics as m
from data import DEFAULT_ROOT, load_bundle, search_symbols
from portfolio import parse_items, parse_weights, portfolio_returns
from portfolios import Portfolio, delete as pf_delete, load_all as pf_load, save as pf_save
from tax import after_tax_portfolio
st.set_page_config(page_title="Stock & Portfolio Analyzer", layout="wide")
# ---------------------------------------------------------------- settings
# user inputs persist across reloads / server restarts / devices
SETTINGS_FILE = Path(__file__).parent / "settings.json"
_settings: dict = {}
try:
_settings = json.loads(SETTINGS_FILE.read_text())
except Exception:
pass
def _remember(**kw) -> None:
"""Persist current widget values (write only on change)."""
global _settings
changed = False
for k, v in kw.items():
if _settings.get(k) != v:
_settings[k] = v
changed = True
if changed:
try:
SETTINGS_FILE.write_text(json.dumps(_settings, indent=2))
except Exception:
pass
# ---------------------------------------------------------------- sidebar
st.sidebar.title("Stock & Portfolio Analyzer")
root = st.sidebar.text_input("Data root", str(DEFAULT_ROOT))
rebuild = st.sidebar.checkbox("Rebuild parquet cache", value=False)
with st.spinner("Loading data..."):
bundle = load_bundle(root=Path(root), rebuild=rebuild)
# data symbols are stored lowercase; input is resolved case-insensitively
_LOWER2SYM = {s.lower(): s for s in bundle.adj.columns}
def _resolve_weights(w: dict) -> dict:
"""Map parsed symbols onto the data's symbol case (insensitive match)."""
return {_LOWER2SYM.get(s.lower(), s): v for s, v in w.items()}
def _replace_token(spec: str, old: str, new: str) -> str:
"""Replace the symbol entry `old` (with or without a ':weight') by `new`.
Grammar: spaces separate items, commas join elements of one item.
"""
def fix_item(item: str) -> str:
out = []
for p in item.split(","):
sym = p.rsplit(":", 1)[0]
if sym.lower() == old.lower():
p = new if ":" not in p else new + p.rsplit(":", 1)[1]
out.append(p)
return ",".join(out)
return " ".join(fix_item(it) for it in spec.split())
st.sidebar.subheader("Portfolio")
# hand-off from e.g. 'Load saved into field': must be applied before the
# spec widget is (re)instantiated
_handover = st.session_state.pop("spec_handover", None)
if _handover is not None:
st.session_state["spec"] = _handover
spec = st.sidebar.text_input(
"Symbol(s) or portfolio(s)", value=_settings.get("spec", ""), key="spec",
help=("One or more, space-separated — same grammar as the Benchmark "
"field.\n"
"MSFT <- single symbol\n"
"MSFT:0.6,V:0.4 <- portfolio (commas join its elements, bare "
"symbols are equal-weight)\n"
"MSFT V googl:0.5,amzn:0.5 <- three separate entries, each "
"analyzed on its own\n"
"The FIRST entry is the 'current' portfolio: it drives the tax "
"detail tab, the allocation tab and Save."))
saved = pf_load()
if saved:
st.sidebar.subheader("Saved portfolios")
lcol1, lcol2 = st.sidebar.columns([3, 1])
load_name = lcol1.selectbox("Load into field", [p.name for p in saved], key="load_name")
if lcol2.button("Load"):
s = next(p.spec for p in saved if p.name == load_name)
# old specs used ', ' between elements; that would now read as two items
st.session_state["spec_handover"] = s.replace(", ", ",")
st.rerun()
_scheme_labels = ["Buy & hold (drift)", "Weekly", "Monthly", "Quarterly", "Yearly"]
scheme = st.sidebar.selectbox(
"Rebalance scheme", _scheme_labels,
index=min(int(_settings.get("scheme_index", 3)), len(_scheme_labels) - 1))
freq = {
"Buy & hold (drift)": None, "Weekly": "1W", "Monthly": "1ME",
"Quarterly": "QE", "Yearly": "YE",
}[scheme]
cost_bps = st.sidebar.number_input("Trading cost (bps, one-way)", 0.0, 100.0,
float(_settings.get("cost_bps", 5.0)))
st.sidebar.subheader("Tax (after-tax analysis)")
lt_rate = st.sidebar.number_input("Long-term gains %", 0.0, 49.0,
float(_settings.get("lt_rate", 20.0))) / 100
st_rate = st.sidebar.number_input("Short-term gains %", 0.0, 49.0,
float(_settings.get("st_rate", 15.0))) / 100
div_rate = st.sidebar.number_input("Dividends %", 0.0, 49.0,
float(_settings.get("div_rate", 15.0))) / 100
st.sidebar.subheader("Period & benchmark")
_p = int(_settings.get("period", 2010))
period = st.sidebar.selectbox("From year", range(2002, 2026),
index=(min(max(_p, 2002), 2025) - 2002))
bench_spec = st.sidebar.text_input(
"Benchmark (optional)", value=_settings.get("bench_spec", ""), key="bench_spec",
help=("Same grammar as the symbol field: a single symbol (MSFT) or a "
"portfolio (MSFT:0.6,V:0.4 — commas join its elements).\n"
"Spaces separate DISTINCT benchmarks, e.g.\n"
"MSFT MSFT:0.6,V:0.4\n"
"Empty = none. Simulated with the same rebalance scheme, cost and "
"tax rules. Beta/alpha in the stats table refers to the first "
"benchmark."))
# remember the sidebar inputs now, before any validation st.stop()
_remember(spec=spec, bench_spec=bench_spec, scheme_index=_scheme_labels.index(scheme),
cost_bps=cost_bps, lt_rate=lt_rate * 100, st_rate=st_rate * 100,
div_rate=div_rate * 100, period=period)
if not spec.strip():
st.info("Type symbols or portfolios in the sidebar — e.g. `MSFT` or "
"`MSFT:0.6,V:0.4`; space-separated entries are each analyzed "
"separately. The page updates as soon as the input is valid.")
st.stop()
try:
items = [_resolve_weights(w) for w in parse_items(spec)]
except ValueError as e:
st.error(f"Invalid input: {e}")
st.stop()
weights = items[0] # first entry = the 'current' portfolio
all_syms = {s for w in items for s in w}
unknown = [s for s in all_syms if s not in bundle.adj.columns]
if unknown:
st.error(f"Symbol field — not in the data: {', '.join(unknown)}")
for u in unknown[:2]:
cands = [s for s in search_symbols(bundle, u, limit=4) if s not in all_syms]
if cands:
st.caption(f"Did you mean {u}:")
cols = st.sidebar.columns(2)
for i, s in enumerate(cands):
if cols[i % 2].button(s, key=f"fix-{u}-{s}", use_container_width=True):
st.session_state["spec_handover"] = _replace_token(spec, u, s)
st.rerun()
st.stop()
start = f"{period}-01-01"
use_syms = list(weights)
# ---- save / delete / select saved portfolios -------------------------
st.sidebar.subheader("Saved portfolios")
if len(items) == 1:
pname = st.sidebar.text_input("Name to save current portfolio as", key="save_name")
if st.sidebar.button("Save current portfolio"):
pname = pname.strip()
if pname:
# commas without spaces: under the spec grammar that stays ONE portfolio
spec_str = ",".join(f"{s}:{w:g}" for s, w in weights.items())
pf_save(Portfolio(pname, spec_str, freq, cost_bps))
st.sidebar.success(f"Saved '{pname}'.")
st.rerun()
else:
st.sidebar.warning("Give it a name first (field below).")
else:
st.sidebar.caption("Save: the Symbol field must hold exactly ONE "
"portfolio to save (it would save the first entry).")
compare_names: list[str] = []
if saved:
compare_names = st.sidebar.multiselect(
"Compare on the chart", [p.name for p in saved],
default=[p.name for p in saved])
dcol1, dcol2 = st.sidebar.columns([3, 1])
del_name = dcol1.text_input("Delete by name", key="del_name")
if dcol2.button("Delete") and del_name.strip():
pf_delete(del_name.strip())
st.rerun()
# ------------------------------------------------------------ compute
# cache the heavy simulations: switching the chart window / curve toggle / tax
# rates only rebuilds the (cheap) HTML — the portfolio+tax runs happen once
@st.cache_data(show_spinner=False)
def _compute_portfolio(w_key: tuple, scheme: str | None, cost_bps: float,
start: str, lt: float, st_r: float, div: float,
root: str):
w = dict(w_key)
r = portfolio_returns(bundle.adj, w, rebalance=scheme, cost_bps=cost_bps, start=start)
t = after_tax_portfolio(bundle.adj, bundle.div, bundle.capg, w, rebalance=scheme,
cost_bps=cost_bps, lt_rate=lt, st_rate=st_r,
div_rate=div, start=start)
return r, t
def build_result(name: str, w: dict, scheme: str | None, cost: float) -> dict:
r, t = _compute_portfolio(tuple(sorted(w.items())), scheme, cost, start,
lt_rate, st_rate, div_rate, str(root))
return {"name": name, "weights": w, "res": r, "tax": t}
results = []
for k, w in enumerate(items):
name = "Current" if len(items) == 1 else ", ".join(w)
results.append(build_result(name, w, freq, cost_bps))
for p in saved:
if p.name not in compare_names:
continue
try:
w = {s: v for s, v in parse_weights(p.spec).items() if s in bundle.adj.columns}
if not w:
st.sidebar.warning(f"'{p.name}': no known symbols, skipped.")
continue
results.append(build_result(p.name, w, p.scheme, p.cost_bps))
except Exception as e:
st.sidebar.warning(f"'{p.name}': {e}")
res, taxres = results[0]["res"], results[0]["tax"]
# benchmark: same spec format as the main field, simulated with the same
# rebalance scheme and cost; invalid/unknown input just disables it
# benchmarks: ';' separates independent specs; each is simulated exactly
# like the candidate portfolio (same scheme, cost and tax rules)
# each item (line, or space-separated on a line) is one benchmark
benchmarks: list[dict] = []
for i, item in enumerate(bench_spec.split(), 1):
try:
w = _resolve_weights(parse_weights(item))
except ValueError as e:
st.sidebar.warning(f"Benchmark {i} ignored: {e}")
continue
unknown = [s for s in w if s not in bundle.adj.columns]
if unknown:
st.sidebar.warning(f"Benchmark {i} ignored: not in the data: "
f"{', '.join(unknown)}")
continue
try:
r, t = _compute_portfolio(tuple(sorted(w.items())), freq, cost_bps, start,
lt_rate, st_rate, div_rate, str(root))
except ValueError as e:
st.sidebar.warning(f"Benchmark {i} ignored: {e}")
continue
benchmarks.append({
"label": ", ".join(w),
"price": r.equity.reindex(res.equity.index).ffill(),
"after": t.equity.reindex(res.equity.index).ffill(),
})
# first benchmark is the reference for beta/alpha in the stats table
bench_price = benchmarks[0]["price"] if benchmarks else None
bench_after = benchmarks[0]["after"] if benchmarks else None
bench_label = benchmarks[0]["label"] if benchmarks else None
# ------------------------------------------------------------ display
if len(items) == 1:
if len(use_syms) == 1:
st.title(f"{use_syms[0]} — single symbol")
else:
st.title(f"Portfolio: {', '.join(use_syms)}")
else:
st.title("Portfolios: " + " · ".join(", ".join(w) for w in items))
if len(results) > 1:
st.caption(f"Comparing: {', '.join(r['name'] for r in results)}")
st.caption(f"{scheme} · start {start} · cost {cost_bps} bps · "
f"tax LT/ST/div {lt_rate:.0%}/{st_rate:.0%}/{div_rate:.0%}"
+ (f" · benchmark: {' ; '.join(b['label'] for b in benchmarks)}"
if benchmarks else ""))
# global curve mode — applies to the stats table AND the chart
_mode_labels = ["Pre-tax", "After-tax", "Pre-tax + after-tax"]
mode = st.radio("Curve", _mode_labels,
horizontal=True, key="curve_mode",
index=_mode_labels.index(_settings["curve_mode"])
if _settings.get("curve_mode") in _mode_labels else 0)
show_pre, show_after, both = mode != "After-tax", mode != "Pre-tax", \
mode == "Pre-tax + after-tax"
tab_stats, tab_equity, tab_alloc, tab_tax = st.tabs(
["Statistics", "Equity curves", "Allocation", "Tax detail"])
with tab_stats:
# one beta/alpha/bench-return column SET per benchmark
bench_ms = [(b["label"], b["price"].resample("ME").last()) for b in benchmarks]
def summarize(price):
out = m.summary(price, None)
for blab, bm in bench_ms:
d = m.summary(price, bm)
sfx = "" if len(bench_ms) == 1 else f" [{blab}]"
out[f"beta{sfx}"] = d["beta"]
out[f"alpha_ann{sfx}"] = d["alpha_ann"]
out[f"ann_return_bench{sfx}"] = d["ann_return_bench"]
return out
# one row per candidate (the whole portfolio, not its components) and
# per benchmark; the pre/after suffix only appears in 'both' mode
summaries = {}
def row(name, pre, after):
if both:
summaries[f"{name} (pre-tax)"] = summarize(pre)
summaries[f"{name} (after-tax)"] = summarize(after)
elif show_pre:
summaries[name] = summarize(pre)
else:
summaries[name] = summarize(after)
for r in results:
row(r["name"], r["res"].equity, r["tax"].equity)
for b in benchmarks:
row(f"benchmark: {b['label']}", b["price"], b["after"])
st.dataframe(m.format_summary_table(summaries), width='stretch')
with tab_equity:
# Self-contained plotly.js page (see chart_widget.py): full mouse zoom +
# pan, with every relayout instantly clamping x to the data and re-fitting
# y to the exact min/max of the visible data. No clipping, no blank space.
from chart_widget import equity_chart_html
_win_labels = ["Max", "10Y", "5Y", "3Y", "1Y"]
win = st.radio("Window", _win_labels,
horizontal=True, key="equity_window",
index=_win_labels.index(_settings["equity_window"])
if _settings.get("equity_window") in _win_labels else 0)
_remember(curve_mode=mode, equity_window=win)
idx = res.equity.index
if win != "Max":
idx = idx[idx >= idx[-1] - pd.DateOffset(years=int(win[:-1]))]
# labels carry the pre/after suffix only in 'both' mode
palette = ["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728",
"#9467bd", "#17becf", "#e377c2", "#8c564b"]
def lab(name, which):
return f"{name}{which}" if both else name
series: dict[str, pd.Series] = {}
styles: dict[str, dict] = {}
for k, r in enumerate(results):
color = palette[k % len(palette)]
if show_pre:
series[lab(r["name"], "pre-tax")] = r["res"].equity
styles[lab(r["name"], "pre-tax")] = {"color": color, "width": 2}
if show_after:
s = lab(r["name"], "after-tax")
series[s] = r["tax"].equity
# faded only in 'both' mode, to stay distinct from the pre-tax line
styles[s] = {"color": color, "width": 2,
"opacity": 0.55 if both else 1.0}
for b in benchmarks:
base = f"benchmark ({b['label']})"
if show_pre:
s = lab(base, "pre-tax")
series[s] = b["price"] / b["price"].iloc[0]
styles[s] = {"width": 1, "opacity": 0.6}
if show_after:
s = lab(base, "after-tax")
series[s] = b["after"] / b["after"].iloc[0]
styles[s] = {"width": 1, "opacity": 0.3 if both else 0.6}
html = equity_chart_html(series, idx, ytitle="growth (1.0 = start)", styles=styles)
st.iframe(html, height=600)
def final(name, pre, after):
if both:
return f"{name}: pre {pre:.2f}× / after {after:.2f}×"
return f"{name}: {(pre if show_pre else after):.2f}×"
finals = " · ".join(final(r["name"], r["res"].equity.iloc[-1],
r["tax"].equity.iloc[-1]) for r in results)
for b in benchmarks:
finals += " · " + final(f"benchmark ({b['label']})", b["price"].iloc[-1],
b["after"].iloc[-1])
st.caption(f"Final: {finals}. "
"Drag to box-zoom, right-drag (or mode-bar hand) to pan, "
"scroll to zoom, double-click to reset.")
with tab_alloc:
st.caption(f"Allocation for the **{results[0]['name']}** portfolio "
"(first entry in the Symbol field).")
alloc = res.allocation
st.dataframe(alloc.tail(1).T, width='stretch')
fig = go.Figure(go.Bar(x=alloc.columns, y=alloc.iloc[-1]))
fig.update_layout(height=360, title=f"Drifted allocation @ {alloc.index[-1].date()}")
st.plotly_chart(fig, width='stretch')
st.caption(f"Final allocation: " +
", ".join(f"{k} {v:.1%}" for k, v in alloc.iloc[-1].items()))
with tab_tax:
st.caption(f"Tax detail for the **{results[0]['name']}** portfolio "
"(first entry in the Symbol field).")
# all values below are fractions of the starting account (1.0 = 100%)
yr_tax = taxres.taxes.resample("YE").sum()
yr_tax.index = yr_tax.index.year
yr_tax = yr_tax.map(lambda v: f"{v:.2%}")
yr_real = taxres.realized.resample("YE").sum()
yr_real.index = yr_real.index.year
yr_real = yr_real.map(lambda v: f"{v:+.2%}")
yr_liq = taxres.liq_tax.resample("YE").last()
yr_liq.index = yr_liq.index.year
yr_liq = yr_liq.map(lambda v: f"{v:.2%}")
st.subheader("Taxes actually paid per year (distributions + rebalance sales)")
st.dataframe(yr_tax, width='stretch')
st.subheader("Realized gains/losses per year (at rebalances)")
st.dataframe(yr_real, width='stretch')
st.subheader("Unrealized-gain tax you would owe if you sold everything (year-end)")
st.dataframe(yr_liq, width='stretch')
st.caption(f"Total taxes paid to date: {taxres.taxes['total'].sum():.2%} of start · "
f"final after-tax value (sell everything): {taxres.equity.iloc[-1]:.2f}× "
f"= {taxres.equity.iloc[-1] - 1:+.1%} total return")

235
chart_widget.py Normal file
View File

@ -0,0 +1,235 @@
"""Interactive equity chart with re-basing zoom.
Streamlit's st.plotly_chart cannot run JS on relayout (scripts are sanitized
and component iframes can't reach the parent chart), so this builds a
self-contained HTML page plotly.js + inline data + zoom logic and embeds
it with st.components.v1.html / st.iframe.
Semantics (the point of this chart):
* The x-axis is clamped to the data's first/last timestamp — you can never
view empty space before or after the data.
* On EVERY view change (zoom, pan, scroll, reset) each series is
RE-BASED: its leftmost visible point is scaled to exactly 1.0.
So whatever window you look at, every line starts at 1.0 on the left
and the chart shows the change in value from that starting point.
* The y-axis is fitted exactly to the re-based visible data no blank
space top or bottom (1.0 is always in view since every line starts there).
"""
from __future__ import annotations
import json
from pathlib import Path
import pandas as pd
_PLOTLY_URL = "https://cdn.plot.ly/plotly-2.35.2.min.js"
_JS_CACHE = Path(__file__).parent / ".cache" / "plotly.min.js"
def _plotly_js_tag() -> str:
"""CDN by default: the browser caches the ~4.6 MB library after the first
load, so chart reloads (window/curve toggles) send a small HTML payload.
Set F_INLINE_PLOTLY=1 for fully offline use (inlines the local copy)."""
import os
if os.environ.get("F_INLINE_PLOTLY"):
if not _JS_CACHE.exists():
_JS_CACHE.parent.mkdir(parents=True, exist_ok=True)
try:
import urllib.request
urllib.request.urlretrieve(_PLOTLY_URL, _JS_CACHE)
except Exception:
pass
if _JS_CACHE.exists():
return f"<script>{_JS_CACHE.read_text()}</script>"
return f'<script src="{_PLOTLY_URL}"></script>'
_TEMPLATE = """<!DOCTYPE html>
<html><head><meta charset="utf-8">{plotly}
<style>html,body{{margin:0;padding:0;background:#fff}}</style>
</head><body>
<div id="c" style="width:100%;height:{height}px"></div>
<script>
// Raw data: growth ratios (1.0 at each series' own start date).
// Raw data: growth ratios (1.0 at each series' own start date), x as ISO dates.
const SERIES = {series_json};
const X0 = SERIES[0].x[0];
const X1 = SERIES.map(s => s.x[s.x.length - 1]).reduce((a, b) => (a > b ? a : b));
const INIT0 = "{init0}", INIT1 = "{init1}";
const ms = v => +new Date(v); // ISO/Date -> epoch ms
const iso = v => new Date(v).toISOString().slice(0, 10);
function lowerBound(a, v) {{ let lo = 0, hi = a.length;
while (lo < hi) {{ const m = (lo + hi) >> 1; if (a[m] < v) lo = m + 1; else hi = m; }}
return lo; }}
function upperBound(a, v) {{ let lo = 0, hi = a.length;
while (lo < hi) {{ const m = (lo + hi) >> 1; if (a[m] <= v) lo = m + 1; else hi = m; }}
return lo; }}
// Re-base every series so its leftmost visible point == 1.0, and compute
// the tight y-range (log10 units) of the visible, re-based data.
function applyView(t0, t1) {{
const norms = SERIES.map(s => {{
const bi = Math.min(lowerBound(s.x, t0), s.x.length - 1);
const base = (s.y[bi] > 0) ? s.y[bi] : 1;
return s.y.map(v => v / base);
}});
let lo = Infinity, hi = -Infinity;
SERIES.forEach((s, k) => {{
const n = norms[k];
for (let i = lowerBound(s.x, t0), j = upperBound(s.x, t1); i < j; i++) {{
if (n[i] < lo) lo = n[i];
if (n[i] > hi) hi = n[i];
}}
}});
if (!isFinite(lo)) {{ lo = 0.99; hi = 1.01; }}
let a = Math.log10(lo), b = Math.log10(hi);
if (b - a < 0.0005) {{ a -= 0.0003; b += 0.0003; }}
return {{ norms, range: [a, b] }};
}}
// First/last VISIBLE data point (ms) across all series within [t0, t1].
function dataBounds(t0, t1) {{
let first = Infinity, last = -Infinity;
for (const s of SERIES) {{
const i = lowerBound(s.x, t0);
if (i < s.x.length) first = Math.min(first, ms(s.x[i]));
const j = upperBound(s.x, t1);
if (j > 0) last = Math.max(last, ms(s.x[j - 1]));
}}
return [first, last];
}}
// X ticks: leftmost/rightmost ALWAYS labeled to the day (YYYY-MM-DD);
// middle ticks get YYYY-MM (or MM-DD for short windows).
function xticks(t0, t1) {{
const m0 = ms(t0), m1 = ms(t1);
const n = 6, vals = [], text = [];
const pad = v => String(v).padStart(2, '0');
const span = m1 - m0;
for (let i = 0; i < n; i++) {{
const d = new Date(m0 + span * i / (n - 1));
vals.push(iso(d));
if (i === 0 || i === n - 1) {{
text.push(d.getUTCFullYear() + '-' + pad(d.getUTCMonth() + 1) + '-' + pad(d.getUTCDate()));
}} else if (span > 62 * 86400000) {{
text.push(d.getUTCFullYear() + '-' + pad(d.getUTCMonth() + 1));
}} else {{
text.push(pad(d.getUTCMonth() + 1) + '-' + pad(d.getUTCDate()));
}}
}}
return {{ vals, text }};
}}
const traces = SERIES.map(s => ({{
x: s.x, y: s.y, type: 'scatter', mode: 'lines', name: s.name, line: s.style
}}));
const gd = document.getElementById('c');
let updating = false;
let lastView = null;
let lastViewMs = null; // [t0, t1] of last rendered view (ms)
function render(t0, t1) {{
// snap the view edges onto the first/last actual data points, so the
// leftmost/rightmost data point sits exactly on the axis edge
const [d0, d1] = dataBounds(t0, t1);
if (isFinite(d0) && isFinite(d1)) {{ t0 = iso(d0); t1 = iso(d1); }}
const {{ norms, range }} = applyView(t0, t1);
const tx = xticks(t0, t1);
lastView = [t0, t1, range[0], range[1]];
lastViewMs = [ms(t0), ms(t1)];
Plotly.restyle(gd, {{ y: norms }}).then(() =>
Plotly.relayout(gd, {{ 'xaxis.range': [t0, t1], 'yaxis.range': range,
'xaxis.tickvals': tx.vals, 'xaxis.ticktext': tx.text }})
).then(() => {{ updating = false; }});
}}
function onRelayout(ev) {{
if (updating) return;
// box-zoom sends xaxis.range[0]/[1]; scroll-zoom/pan send xaxis.range
if (!Object.keys(ev).some(k => k.startsWith('xaxis.range'))) return;
updating = true;
const rA = ms(gd.layout.xaxis.range[0]);
const rB = ms(gd.layout.xaxis.range[1]);
const cA = lastViewMs ? lastViewMs[0] : ms(X0);
const cB = lastViewMs ? lastViewMs[1] : ms(X1);
const minSpan = 7 * 86400000; // 1 week at full zoom
const maxSpan = ms(X1) - ms(X0);
let a, b;
if (Math.abs((rB - rA) - (cB - cA)) < 1e-6 * Math.max(cB - cA, 1)) {{
// PAN: the window keeps its width and sticks to the data edges
// (clamping only one edge would shrink the window -> looked like zoom).
const W = Math.min(Math.max(rB - rA, minSpan), maxSpan);
a = rA; b = rA + W;
if (a < ms(X0)) {{ a = ms(X0); b = a + W; }}
if (b > ms(X1)) {{ b = ms(X1); a = b - W; }}
}} else {{
// ZOOM (scroll/box/reset): clamp the edges to the data span
a = Math.max(rA, ms(X0));
b = Math.min(rB, ms(X1));
if (b - a < minSpan) {{
if (b >= ms(X1)) {{ b = ms(X1); a = Math.max(ms(X0), ms(X1) - minSpan); }}
else if (a <= ms(X0)) {{ a = ms(X0); b = Math.min(ms(X1), ms(X0) + minSpan); }}
else {{ const c = (a + b) / 2; a = c - minSpan / 2; b = c + minSpan / 2; }}
}}
}}
const t0 = iso(a), t1 = iso(b);
// skip only if the *current* view already matches the clamped target
if (lastView && t0 === lastView[0] && t1 === lastView[1]
&& iso(gd.layout.xaxis.range[0]) === t0 && iso(gd.layout.xaxis.range[1]) === t1) {{
updating = false; return; // truly nothing to do
}}
render(t0, t1);
}}
Plotly.newPlot(gd, traces, {{
margin: {{ l: 60, r: 55, t: 10, b: 90 }}, // b: room for x labels + legend below
xaxis: {{ type: 'date', rangeslider: {{ visible: false }},
hoverformat: '%Y-%m-%d' }},
yaxis: {{ type: 'log', title: '{ytitle}' }},
legend: {{ orientation: 'h', y: -0.30 }}, // underneath the plot, not over it
dragmode: 'zoom',
hovermode: 'x unified'
}}, {{
responsive: true, scrollZoom: true, doubleClick: 'reset',
modeBarButtonsToRemove: ['autoScale2d', 'select2d', 'lasso2d', 'toImage']
}}).then(() => {{
gd.on('plotly_relayout', onRelayout);
updating = true;
render(INIT0, INIT1);
}});
</script></body></html>
"""
def equity_chart_html(series: dict[str, pd.Series], window: pd.DatetimeIndex,
height: int = 540, ytitle: str = "growth",
styles: dict[str, dict] | None = None) -> str:
"""Render growth-ratio series as a self-contained interactive HTML page.
series: {name: ratio Series} (each starts at 1.0 on its own first date).
window: initial visible DatetimeIndex (x is always clamped to this span).
"""
styles = styles or {}
payload = []
for name, s in series.items():
s = s.reindex(window).ffill().dropna()
if s.empty:
continue
payload.append({
"name": name,
"x": [t.date().isoformat() for t in s.index],
"y": [float(v) for v in s.to_numpy(dtype=float)],
"style": styles.get(name, {"width": 2}),
})
w = pd.DatetimeIndex(window)
return _TEMPLATE.format(
plotly=_plotly_js_tag(),
height=height,
series_json=json.dumps(payload),
init0=w[0].date().isoformat(),
init1=w[-1].date().isoformat(),
ytitle=ytitle,
)

136
data.py Normal file
View File

@ -0,0 +1,136 @@
"""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
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"
@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_panel(root: Path, suffix: str, col: str) -> pd.DataFrame:
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:
continue
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_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
# 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] = {}
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")]
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
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]

221
families.py Normal file
View File

@ -0,0 +1,221 @@
"""Fund-family classification.
Tags each symbol with its fund family (Vanguard, Fidelity, iShares/BlackRock,
T. Rowe Price, ...) from the Yahoo `longName`/`shortName` in the per-symbol
chart JSON, plus ticker-pattern fallbacks for the big families whose names
are sometimes missing.
Run: python3 -m families -> writes .cache/families.csv (sym, family, name)
"""
from __future__ import annotations
import csv
import json
import re
from collections import Counter
from pathlib import Path
ROOT = Path.home() / "prog/fin/stocks"
OUT = Path(__file__).parent / ".cache" / "families.csv"
# Ordered keyword rules: first match wins. Specific brands before generics.
RULES: list[tuple[str, str]] = [
("Vanguard", r"vanguard"),
("Fidelity", r"fidelity|fd[ae]{1,2}[uiv]?$|fdsi|fdls|fdsm"),
("iShares/BlackRock", r"\bishares\b"),
("BlackRock", r"blackrock"),
("State Street/SPDR", r"state street|\bspdr\b|spts|sx[ae]x$|sfx[ae]$"),
("T. Rowe Price", r"t\.? ?rowe price|^\btr[pd]|tlc[ae]d$|tltlx$|thdax$|trn[ax]$"),
("Schwab", r"schwab"),
("MassMutual", r"massmutual|mass mutual"),
("Franklin", r"franklin"),
("Templeton", r"templeton"),
("J.P. Morgan", r"j\.? ?p\.? ?morgan|jpmorgan|jps[ae]x$|jp[ae]x$"),
("Invesco", r"invesco"),
("American Century", r"american century"),
("Lord Abbett", r"lord abbett"),
("Victory Capital", r"victory capital|victory shares|victory pioneers|victory pioneer|victory rs|victory fund"),
("DFA/Dimensional", r"dimensional|\bdfa\b|df[ae]x$|dgbex$|d[ij]a[ei]x?$"),
("Morgan Stanley", r"morgan stanley"),
("Voya", r"\bvoya\b"),
("Ned Davis", r"ned davis"),
("Neuberger Berman", r"neuberger"),
("Dodge & Cox", r"dodge ?cox"),
("Janus Henderson", r"janus"),
("PIMCO", r"pimco"),
("Northern Trust", r"northern trust"),
("Wells Fargo", r"wells fargo"),
("Barclays", r"barclays"),
("Carillon", r"carillon"),
("Legg Mason", r"legg mason"),
("Allianz", r"allianz"),
("American Funds/Alliance", r"american funds|american fds"),
("American Beacon", r"american beacon|\bab (small|high|equity|income|bond)"),
("BNY Mellon", r"bny |mellon"),
("Goldman Sachs AM", r"goldman sachs"),
("Nuveen", r"nuveen"),
("Western Asset", r"western asset"),
("Loomis Sayles", r"loomis sayles"),
("MFS/Merrill", r"\bmfs\b"),
("Principal", r"principal "),
("PGIM", r"pgim"),
("Hartford", r"hartford"),
("ClearBridge", r"clearbridge"),
("Nationwide", r"nationwide"),
("abrdn", r"abrdn"),
("John Hancock", r"john hancock|jhanco|\bjh[a-z]{2}x$"),
("Virtus", r"virtus"),
("SEI", r"\bsei\b"),
("Empower", r"empower"),
("Avantis", r"avantis"),
("Allspring", r"allspring"),
("Baird", r"baird"),
("Nomura", r"nomura"),
("Baron", r"baron "),
("Artisan", r"artisan"),
("Columbia", r"columbia "),
("Calvert", r"calvert"),
("First Eagle", r"first eagle"),
("AQR", r"\baqr\b"),
("Catalyst", r"catalyst"),
("Touchstone", r"touchstone"),
("NexPoint", r"nexpoint"),
("Affiliated", r"affiliated managers"),
("Putnam", r"putnam"),
("Federated", r"federated"),
("Schroders", r"schroder"),
("Alger", r"alger "),
("Oakmark", r"oakmark"),
("DoubleLine", r"doubleline"),
("Diamond Hill", r"diamond hill"),
("First Trust", r"first trust|firsttrust"),
("Transamerica", r"transamerica"),
("Brookfield", r"brookfield"),
("DWS", r"\bdws\b"),
("Natixis", r"natixis"),
("Barings/MML", r"mml |barings"),
("Redmont", r"redmont"),
("Parametric", r"parametric"),
("Brown Advisory", r"brown advisory"),
("Harbor", r"harbor "),
("Madison", r"madison "),
("Toews", r"toews"),
("Conestoga", r"conestoga"),
("Turner", r"turner funds"),
("Cambiar", r"cambiar"),
("Copeland", r"copeland"),
("Pear Tree", r"pear tree"),
("Iron/Unified", r"\biron\b.*fund|unified series"),
("E-Valuator", r"e-valuator"),
("SilverPepper", r"silverpepper"),
("Global X", r"global x"),
("Eaton Vance", r"eaton vance"),
("William Blair", r"william blair"),
("Xtrackers/DWS", r"xtrackers"),
("DWS", r"\bdws\b"),
("WisdomTree", r"wisdomtree"),
("Innovator", r"innovator "),
("iPath/UBS", r"ipath"),
("UBS", r"\bubs\b"),
("Cohen & Steers", r"cohen & steers|cohen"),
("Royce", r"royce"),
("Pacer", r"pacer "),
("Capital Group", r"capital group"),
("FlexShares", r"flexshares"),
("Alps", r"alps "),
("Segall Bryant", r"segall bryant"),
("Guinness Atkinson", r"guinness"),
("Calamos", r"calamos"),
("Glenmede", r"glenmede"),
("KraneShares", r"kraneshares"),
("ProShares", r"proshares"),
("VanEck", r"vaneck|van eck"),
("Guggenheim", r"guggenheim"),
("Royce", r"royce"),
("Harding Loevner", r"harding"),
("Burnham", r"burnham"),
("Leuthold", r"leuthold"),
("LSV Global", r"lsv "),
("Cambria", r"cambria"),
("QRAFT", r"qraft"),
("TrueShares", r"trueshares"),
("NorthSquare", r"north ?square"),
("AdvisorShares", r"advisors?hares"),
("GQ Partners", r"gqg partners|gq partners"),
("Brandywine", r"brandywine"),
("RBB", r"rbb "),
("Strategy Shares", r"strategy shares"),
("Overlay Shares", r"overlay shares"),
("Robo Global", r"robo global"),
("Henry Schwebel", r"schwebel|amg gw"),
("HC Fund Mgmt", r"hcm "),
("Investment Mgrs Series", r"investment managers series"),
("United States (ETF)", r"^united states |united states (corporate|government|bond)"),
]
_RULES = [(name, re.compile(pat, re.I)) for name, pat in RULES]
# Ticker-pattern fallbacks (used when the name is empty or unmatched).
TICKER_RULES: list[tuple[str, str]] = [
("Vanguard", r"^vt[bcdfghjkmnpqrsuw][a-z]x?$"),
("Fidelity", r"^fd[a-z]{2,4}$"),
("iShares/BlackRock", r"^i[a-z]{3}$"),
("State Street/SPDR", r"^s[pfx][a-z]{2}$|^st[a-z]{2,3}x?$"),
("T. Rowe Price", r"^t[lm][a-z]x?$"),
("MassMutual", r"^m[a-z]{4}$"),
("Franklin", r"^f[ae][a-z]{2,3}x?$"),
("DFA/Dimensional", r"^d[fgij][ae]x?$|^df[a-z]x?$"),
("Schwab", r"^sw[a-z]{2,3}$"),
("Vanguard", r"^vb[a-z]x?$"),
("Global X", r"^[a-z]{3}x$"),
("WisdomTree", r"^w[a-z]{3}$"),
("Xtrackers/DWS", r"^x[a-z]{3}$"),
("Pacer", r"^p[a-z]{2}x$"),
]
_TICKER_RULES = [(name, re.compile(pat, re.I)) for name, pat in TICKER_RULES]
UNKNOWN = "(unknown)"
def classify(sym: str, name: str) -> str:
n = (name or "").strip()
if n:
for fam, rx in _RULES:
if rx.search(n):
return fam
for fam, rx in _TICKER_RULES:
if rx.fullmatch(sym):
return fam
return UNKNOWN
def load_names(root: Path = ROOT) -> dict[str, str]:
out: dict[str, str] = {}
for f in root.glob("*.json"):
sym = f.name.rsplit(".", 1)[0]
try:
m = json.loads(f.read_bytes())["chart"]["result"][0]["meta"]
out[sym] = (m.get("longName") or m.get("shortName") or "").strip()
except Exception:
out[sym] = ""
return out
def build(root: Path = ROOT, out: Path = OUT) -> Counter:
names = load_names(root)
out.parent.mkdir(parents=True, exist_ok=True)
with out.open("w", newline="") as fh:
w = csv.writer(fh)
w.writerow(["sym", "family", "name"])
for sym in sorted(names):
w.writerow([sym, classify(sym, names[sym]), names[sym]])
return Counter(classify(s, n) for s, n in names.items())
if __name__ == "__main__":
c = build()
total = sum(c.values())
print(f"{len(c)} families, {total} symbols")
for k, v in c.most_common(40):
print(f"{v:6d} {k}")
print(f"\nunknown: {c[UNKNOWN]} ({c[UNKNOWN]/total:.1%})")

102
metrics.py Normal file
View File

@ -0,0 +1,102 @@
"""Performance statistics on price/return series (daily, 252 days/yr).
All functions accept a price series (or DataFrame) and return scalars or
Series. Kept deliberately dependency-light (pandas/numpy only) so every
metric is transparent and tweakable. `empyrical-reloaded` is a fine
drop-in for more metrics if you ever want them.
"""
from __future__ import annotations
import numpy as np
import pandas as pd
ANN = 252
def daily_returns(price: pd.Series | pd.DataFrame) -> pd.Series | pd.DataFrame:
return price.pct_change().fillna(0.0)
def total_return(price: pd.Series | pd.DataFrame) -> float:
return float(price.iloc[-1] / price.iloc[0] - 1.0)
def annualized_return(price: pd.Series | pd.DataFrame) -> float:
n = len(price)
return float((price.iloc[-1] / price.iloc[0]) ** (ANN / n) - 1.0)
def annualized_vol(returns: pd.Series | pd.DataFrame) -> float:
return float(returns.std() * np.sqrt(ANN))
def sharpe(returns: pd.Series, rf: float = 0.0) -> float:
r = returns - rf / ANN
sd = r.std()
return float(r.mean() / sd * np.sqrt(ANN)) if sd > 0 else 0.0
def sortino(returns: pd.Series, rf: float = 0.0) -> float:
r = returns - rf / ANN
dd = float(np.sqrt(np.mean(np.minimum(r, 0.0) ** 2)))
return float(r.mean() / dd * np.sqrt(ANN)) if dd > 0 else 0.0
def max_drawdown(price: pd.Series | pd.DataFrame) -> float:
peak = price.cummax()
return float((price / peak - 1.0).min())
def calmar(price: pd.Series | pd.DataFrame) -> float:
mdd = max_drawdown(price)
return float(annualized_return(price) / -mdd) if mdd < 0 else 0.0
def beta_alpha(returns: pd.Series, bench: pd.Series, rf: float = 0.0):
"""CAPM regression. Returns (beta, annualized_alpha)."""
r = (returns - rf / ANN).dropna()
b = (bench - rf / ANN).dropna()
r, b = r.align(b, join="inner")
beta = np.cov(r, b)[0, 1] / np.var(b)
alpha_daily = r.mean() - (rf / ANN + beta * (b.mean() - rf / ANN))
return float(beta), float(alpha_daily * ANN)
def summary(price: pd.Series, bench: pd.Series | None = None,
rf: float = 0.0) -> dict[str, float]:
r = daily_returns(price)
out = {
"total_return": total_return(price),
"ann_return": annualized_return(price),
"ann_vol": annualized_vol(r),
"sharpe": sharpe(r, rf),
"sortino": sortino(r, rf),
"max_dd": max_drawdown(price),
"calmar": calmar(price),
}
if bench is not None:
b, a = beta_alpha(r, daily_returns(bench), rf)
out["beta"] = b
out["alpha_ann"] = a
out["ann_return_bench"] = annualized_return(bench)
return out
def format_summary_table(summaries: dict[str, dict[str, float]]) -> pd.DataFrame:
"""{label: summary_dict} -> transposed table, percentages pre-formatted.
Benchmark columns may carry a ' [<benchmark>] ' suffix when there are
several benchmarks; all of them are formatted by prefix.
"""
df = pd.DataFrame(summaries).T
pct = ("total_return", "ann_return", "ann_vol", "max_dd", "alpha_ann",
"ann_return_bench")
two = ("sharpe", "sortino", "calmar", "beta")
for c in df.columns:
base = c.split(" [")[0]
if base in pct:
df[c] = df[c].map(lambda v: f"{v:,.1%}")
elif base in two:
df[c] = df[c].map(lambda v: f"{v:.2f}")
return df

147
portfolio.py Normal file
View File

@ -0,0 +1,147 @@
"""Portfolio construction: weighted blend of assets with rebalancing.
All inputs are date x symbol DataFrames of *adjusted* prices (which
already include distributions). Returns are pre-tax.
Supported schemes:
rebalance=None buy & hold (drift for the whole period)
rebalance='1W'/'1ME'/'QE'/'YE' etc. (pandas offset alias)
rebalance to target weights on the first trading
day of each period
cost_bps one-way trading cost in basis points, applied on
the traded fraction (turnover) at each rebalance
The weight evolution between rebalances is the standard drift:
w_{t+1} = w_t (1 + r_t) / (1 + R_t)
"""
from __future__ import annotations
import math
from dataclasses import dataclass, field
import numpy as np
import pandas as pd
@dataclass
class PortfolioResult:
equity: pd.Series # gross value, starts at 1.0
returns: pd.Series # daily portfolio return
allocation: pd.DataFrame # drifted weights (date x asset)
turnover: pd.Series # one-way turnover at each rebalance (else 0)
costs: pd.Series # cost paid at each rebalance (else 0)
rebalance_dates: list = field(default_factory=list)
def _rebalance_dates(index: pd.DatetimeIndex, freq: str) -> set:
"""First trading day of each period, including the first day."""
periods = index.to_series().groupby(pd.Grouper(freq=freq)).first()
return set(periods.index.tolist())
def portfolio_returns(adj: pd.DataFrame, weights: dict[str, float],
rebalance: str | None = None,
cost_bps: float = 0.0,
start: str | None = None, end: str | None = None
) -> PortfolioResult:
"""Simulate a (drifting) weighted portfolio on an adjusted price panel.
weights: {symbol: raw weight}; normalized internally.
"""
p = adj[list(weights)].copy()
if start or end:
p = p.loc[start:end]
p = p.dropna()
if len(p) < 2:
raise ValueError("no overlapping data for the given symbols/period")
w0 = np.array([weights[s] for s in p.columns], dtype=float)
w0 = w0 / w0.sum()
rets = p.pct_change().fillna(0.0).values
rebal_set = _rebalance_dates(p.index, rebalance) if rebalance else set()
cost = cost_bps / 1e4
n_days = len(p)
w = w0.copy()
eq = np.empty(n_days)
alloc = np.empty((n_days, len(p.columns)))
turns = np.zeros(n_days)
costs = np.zeros(n_days)
eq[0] = 1.0
alloc[0] = w
# initial buy: cost on 100% invested
eq[0] *= (1.0 - cost * 1.0)
costs[0] = cost
for i in range(1, n_days):
R = float(np.dot(w, rets[i]))
w_drifted = w * (1.0 + rets[i]) / (1.0 + R)
if p.index[i] in rebal_set:
target = w0
turnover = float(np.abs(target - w_drifted).sum() / 2.0)
c = cost * turnover
w_drifted = target
eq[i] = eq[i - 1] * (1.0 + R) * (1.0 - c)
turns[i] = turnover
costs[i] = c
else:
eq[i] = eq[i - 1] * (1.0 + R)
w = w_drifted
alloc[i] = w
idx = p.index
return PortfolioResult(
equity=pd.Series(eq, index=idx),
returns=pd.Series(pd.Series(eq, index=idx).pct_change().fillna(0.0), index=idx),
allocation=pd.DataFrame(alloc, index=idx, columns=p.columns),
turnover=pd.Series(turns, index=idx),
costs=pd.Series(costs, index=idx),
rebalance_dates=sorted(rebal_set & set(idx)),
)
def parse_weights(spec: str) -> dict[str, float]:
"""Parse ONE symbol or ONE portfolio spec. Returns {symbol: weight}.
Commas separate the elements of the portfolio; each element is
'SYM' (equal weight) or 'SYM:w' (weight w > 0). A bare 'SYM' on its
own is a single symbol, not a portfolio. Surrounding whitespace is
ignored. For specs that may hold several space-separated
symbols/portfolios, use parse_items().
"""
parts = [t.strip() for t in spec.strip().split(",") if t.strip()]
if not parts:
raise ValueError("no symbols given")
out: dict[str, float] = {}
for token in parts:
if ":" in token:
sym, w_s = token.rsplit(":", 1)
try:
w = float(w_s)
except ValueError:
raise ValueError(f"bad weight in '{token}'") from None
if not math.isfinite(w) or w <= 0:
raise ValueError(f"weight in '{token}' must be a positive number")
else:
sym, w = token, 1.0
sym = sym.strip()
if not sym:
raise ValueError(f"empty symbol in '{token}'")
out[sym] = out.get(sym, 0.0) + w
if not out:
raise ValueError("no symbols parsed")
return out
def parse_items(spec: str) -> list[dict[str, float]]:
"""Parse a spec that may hold several symbols/portfolios.
Spaces (or newlines) separate DISTINCT symbols/portfolios; commas
join the elements of one portfolio. Each item is parsed by
parse_weights().
"""
return [parse_weights(item) for item in spec.split() if item.strip()]

1
portfolios.json Normal file
View File

@ -0,0 +1 @@
[]

42
portfolios.py Normal file
View File

@ -0,0 +1,42 @@
"""Saved portfolio definitions, persisted to portfolios.json next to the app.
A portfolio is fully defined by its symbol/weight spec, rebalance scheme
and cost. Tax rates and the analysis period are global (sidebar), not part
of the saved definition.
"""
from __future__ import annotations
import json
from dataclasses import asdict, dataclass
from pathlib import Path
FILE = Path(__file__).parent / "portfolios.json"
@dataclass
class Portfolio:
name: str
spec: str # "sym1:0.5, sym2:0.5" (parsed by portfolio.parse_weights)
scheme: str | None # rebalance freq alias ("1W", "1ME", "QE", "YE") or None
cost_bps: float = 0.0
def load_all() -> list[Portfolio]:
if not FILE.exists():
return []
try:
return [Portfolio(**d) for d in json.loads(FILE.read_text())]
except Exception:
return []
def save(p: Portfolio) -> None:
all_p = {q.name: q for q in load_all()}
all_p[p.name] = p
FILE.write_text(json.dumps([asdict(q) for q in all_p.values()], indent=2))
def delete(name: str) -> None:
all_p = [q for q in load_all() if q.name != name]
FILE.write_text(json.dumps([asdict(q) for q in all_p], indent=2))

7
requirements.txt Normal file
View File

@ -0,0 +1,7 @@
pandas>=2.0
numpy
pyarrow>=14
plotly>=5.18
streamlit>=1.32
# optional: more off-the-shelf performance metrics
# empyrical-reloaded

15
run.sh Executable file
View File

@ -0,0 +1,15 @@
#!/usr/bin/env bash
# Launch the Stock & Portfolio Analyzer web UI.
# Fixed port: 8599 (http://localhost:8599)
set -euo pipefail
cd "$(dirname "$0")"
if pgrep -f "streamlit run app\.py" >/dev/null; then
echo "already running: $(pgrep -af 'streamlit run app\.py' | head -1)"
exit 0
fi
exec .venv/bin/streamlit run app.py \
--server.port 8599 \
--server.headless true \
--browser.gatherUsageStats false

25
run_tests.sh Executable file
View File

@ -0,0 +1,25 @@
#!/usr/bin/env bash
# Full test suite: app-level (Streamlit AppTest) + real-browser e2e (Playwright).
#
# Usage: ./run_tests.sh
# The browser e2e part needs the server running (./run.sh, port 8599);
# it is skipped automatically if the server is down.
set -uo pipefail
cd "$(dirname "$0")"
fail=0
echo "=== 1/2 app tests (AppTest, no browser) ==="
.venv/bin/python tests/test_app.py || fail=1
echo
echo "=== 2/2 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
echo "server not reachable on :8599 — skipping e2e (start it with ./run.sh)"
fi
echo
[ $fail -eq 0 ] && echo "ALL TESTS PASSED" || echo "TESTS FAILED"
exit $fail

200
tax.py Normal file
View File

@ -0,0 +1,200 @@
"""After-tax portfolio value: "what do I keep if I sell everything today?"
Model (simplified DAS)
----------------------
The account starts at 1.0 (growth-ratio units; no fixed capital).
- Holdings are valued on *adjusted* prices, which already assume
distributions are reinvested. Distributions therefore flow through as:
dividend income -> taxed at `div_rate`
capital gain dist-> taxed at `lt_rate`
The after-tax remainder increases each lot's cost basis proportionally.
- On rebalance, sells are FIFO. A lot is long-term if held more than
365 days at sale, else short-term; realized gains/losses are taxed at
`lt_rate` / `st_rate`.
- THE headline number, `equity[t]`, is the after-tax value if you
LIQUIDATE the entire account on day t:
equity[t] = market value + cash - tax_on_unrealized_gains(t)
where tax_on_unrealized_gains marks every remaining lot to market,
classifies it LT/ST by age, nets losses against gains, and applies the
rates. This is what you would actually have in your pocket after
selling everything and filing your taxes.
Deliberately NOT modeled: loss carryover, wash sales, state rates,
brackets.
"""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
import pandas as pd
ST_WINDOW_DAYS = 365
@dataclass
class _Lot:
units: float
cost: float
date: pd.Timestamp
@dataclass
class TaxResult:
equity: pd.Series # after-tax value if you sell everything today (ratio)
liq_tax: pd.Series # tax that a full liquidation today would owe (unrealized)
taxes: pd.DataFrame # daily taxes actually PAID: div_tax, capg_tax, realized_tax, total
realized: pd.DataFrame # daily realized (at rebalances): lt_gain, st_gain, lt_loss, st_loss
lots_outstanding: int # final lot count (sanity check)
def _liquidation_tax(lt_gain: float, lt_loss: float,
st_gain: float, st_loss: float,
lt_rate: float, st_rate: float) -> float:
"""Tax on selling all lots now; losses offset gains (same/other type)."""
lt_net = lt_gain - lt_loss
st_net = st_gain - st_loss
if lt_net >= 0 and st_net >= 0:
return lt_net * lt_rate + st_net * st_rate
if lt_net < 0 and st_net < 0:
return 0.0
if st_net < 0: # ST losses offset LT gains
return max(lt_net + st_net, 0.0) * lt_rate
return max(st_net + lt_net, 0.0) * st_rate
def after_tax_portfolio(adj: pd.DataFrame, div: pd.DataFrame, capg: pd.DataFrame,
weights: dict[str, float],
rebalance: str | None = None, cost_bps: float = 0.0,
lt_rate: float = 0.20, st_rate: float = 0.15,
div_rate: float = 0.15,
start: str | None = None, end: str | None = None
) -> TaxResult:
syms = [s for s in weights if s in adj.columns]
p = adj[syms]
d = div[[s for s in syms if s in div.columns]].reindex(index=p.index, columns=syms).fillna(0.0)
c = capg[[s for s in syms if s in capg.columns]].reindex(index=p.index, columns=syms).fillna(0.0)
if start or end:
p = p.loc[start:end]
p = p.dropna()
d, c = d.reindex(p.index).fillna(0.0), c.reindex(p.index).fillna(0.0)
if len(p) < 2:
raise ValueError("no overlapping data for the given symbols/period")
idx = p.index
pv = p.values
dv = d.values
cv = c.values
n = len(p)
w0 = np.array([weights[s] for s in syms])
w0 = w0 / w0.sum()
cost = cost_bps / 1e4
# state (account = 1.0 at start)
units = w0 / pv[0]
lots: list[list[_Lot]] = [[_Lot(units[i], float(units[i] * pv[0][i]), idx[0])]
if units[i] > 0 else [] for i in range(len(syms))]
cash = -float(w0.sum() * cost) # initial purchase cost
tax_cols = ["div_tax", "capg_tax", "realized_tax", "total"]
tax_rows = np.zeros((n, len(tax_cols)))
real_cols = ["lt_gain", "st_gain", "lt_loss", "st_loss"]
real_rows = np.zeros((n, len(real_cols)))
equity = np.empty(n)
liq_tax = np.zeros(n)
# rebalance schedule (skip first day)
from portfolio import _rebalance_dates
rebal = (_rebalance_dates(idx, rebalance) - {idx[0]}) if rebalance else set()
for t in range(n):
prices = pv[t]
market_value = float(np.dot(units, prices))
# --- distributions (taxed, net flows back into basis) -----------
dinc = units * dv[t]
cinc = units * cv[t]
d_tax = float(dinc.sum() * div_rate)
c_tax = float(cinc.sum() * lt_rate)
cash -= d_tax + c_tax
tax_rows[t, 0] = d_tax
tax_rows[t, 1] = c_tax
for i in range(len(syms)): # grow cost basis with reinvested net
net = (dv[t, i] * (1 - div_rate) + cv[t, i] * (1 - lt_rate)) * units[i]
if net <= 0 or not lots[i]:
continue
tot = sum(l.units for l in lots[i])
for l in lots[i]:
l.cost += net * (l.units / tot)
# --- rebalance to target weights --------------------------------
if idx[t] in rebal:
value = market_value
target_val = w0 * value
for i in range(len(syms)):
cur_val = units[i] * prices[i]
trade = target_val[i] - cur_val # + buy, - sell
if abs(trade) < 1e-9:
continue
c_cost = cost * abs(trade)
cash -= c_cost
if trade < 0: # sell |trade| at price, FIFO
to_sell = -trade / prices[i]
for l in lots[i]:
if to_sell <= 1e-12:
break
u = min(l.units, to_sell)
gain = u * prices[i] - (l.cost * u / l.units)
held = (idx[t] - l.date).days
li = 1 if held > ST_WINDOW_DAYS else 2
if gain >= 0:
tax = gain * (lt_rate if li == 1 else st_rate)
real_rows[t, li - 1] += gain
else:
tax = gain * (lt_rate if li == 1 else st_rate)
real_rows[t, 3 if li == 1 else 2] += -gain
cash -= tax
tax_rows[t, 2] += tax
l.units -= u
l.cost *= (l.units / (l.units + u)) if l.units > 0 else 0.0
to_sell -= u
lots[i] = [l for l in lots[i] if l.units > 1e-12]
units[i] = cur_val / prices[i] + trade / prices[i]
else: # buy
spend = trade + c_cost
u = trade / prices[i]
lots[i].append(_Lot(u, spend, idx[t]))
units[i] += u
# --- tax owed if we liquidate everything today ------------------
lt_gain = lt_loss = st_gain = st_loss = 0.0
for i in range(len(syms)):
px = prices[i]
for l in lots[i]:
gain = l.units * px - l.cost
if (idx[t] - l.date).days > ST_WINDOW_DAYS:
if gain >= 0:
lt_gain += gain
else:
lt_loss += -gain
else:
if gain >= 0:
st_gain += gain
else:
st_loss += -gain
liq_tax[t] = _liquidation_tax(lt_gain, lt_loss, st_gain, st_loss,
lt_rate, st_rate)
equity[t] = market_value + cash - liq_tax[t]
tax_rows[t, 3] = tax_rows[t, 0] + tax_rows[t, 1] + tax_rows[t, 2]
return TaxResult(
equity=pd.Series(equity, index=idx, name="after_tax_liquidation"),
liq_tax=pd.Series(liq_tax, index=idx, name="liquidation_tax"),
taxes=pd.DataFrame(tax_rows, index=idx, columns=tax_cols),
realized=pd.DataFrame(real_rows, index=idx, columns=real_cols),
lots_outstanding=sum(len(l) for l in lots),
)

178
tests/test_app.py Normal file
View File

@ -0,0 +1,178 @@
"""App-level tests (Streamlit AppTest — no browser needed).
Run: .venv/bin/python tests/test_app.py
Exit code 0 = all passed.
Memory note: one data-bundle load is ~2.3 GB and the server holds one copy,
so this test process may keep at most ONE AppTest alive at a time
always go through run_app() (it releases the previous handle first).
"""
from __future__ import annotations
import gc
import json
import logging
import pathlib
import sys
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))
logging.getLogger("streamlit").setLevel(logging.ERROR)
from streamlit.testing.v1 import AppTest # noqa: E402
import portfolios as pf # noqa: E402
ROOT = pathlib.Path(__file__).resolve().parent.parent
SETTINGS = ROOT / "settings.json"
PASS, FAIL = 0, 0
_A: AppTest | None = None
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 app() -> AppTest:
"""The single live AppTest (set by run_app)."""
return _A
def run_app(spec: str | None = None, bench: str | None = None,
**extra_state) -> AppTest:
"""Run the app with the given committed inputs.
spec/bench=None leaves the widget default (restored from settings.json).
Releases the previous AppTest first to bound memory.
"""
global _A
_A = None
gc.collect()
_A = AppTest.from_file(str(ROOT / "app.py"), default_timeout=120)
if spec is not None:
_A.session_state["spec"] = spec
if bench is not None:
_A.session_state["bench_spec"] = bench
for k, v in extra_state.items():
_A.session_state[k] = v
_A.run()
return _A
def main() -> int:
saved_settings = SETTINGS.read_text() if SETTINGS.exists() else None
SETTINGS.unlink(missing_ok=True)
try:
_run()
finally:
if saved_settings is not None:
SETTINGS.write_text(saved_settings)
for name in ("ttest", "oldfmt"):
pf.delete(name)
print(f"\n{PASS} passed, {FAIL} failed")
return 1 if FAIL else 0
def _run() -> None:
# NOTE: never bind app() to a local that outlives the next run_app() —
# each live AppTest pins a ~2.3 GB bundle; there may be at most one.
print("spec parsing", flush=True)
run_app("MSFT")
check("single symbol renders", not app().exception and
any("single symbol" in t.value for t in app().title))
run_app("MSFT:0.6,V:0.4")
check("comma portfolio renders", not app().exception and
any(t.value == "Portfolio: msft, v" for t in app().title))
run_app("MSFT V googl:0.5,amzn:0.5")
check("multiple entries render", not app().exception and
any("Portfolios:" in t.value for t in app().title))
run_app("MSFT:xyz")
check("bad weight -> error", not app().exception and
any("Invalid input" in e.value for e in app().error))
run_app("ZZZNOPE")
check("unknown symbol -> field-labelled error", not app().exception and
any("Symbol field" in e.value and "not in the data" in e.value
for e in app().error))
run_app("")
check("empty -> info, no crash", not app().exception and bool(app().info))
print("benchmarks", flush=True)
run_app("MSFT", "V googl:0.5,amzn:0.5")
tbl = app().main.tabs[0].dataframe[0].value
check("two benchmarks, one line", not app().exception and
"benchmark: v" in list(tbl.index) and
"benchmark: googl, amzn" in list(tbl.index))
check("per-benchmark beta columns",
any(c.startswith("beta [v]") for c in tbl.columns) and
any(c.startswith("beta [googl, amzn]") for c in tbl.columns))
run_app("MSFT", "V zzzznope")
rows = list(app().main.tabs[0].dataframe[0].value.index)
check("invalid benchmark warns, valid survives", not app().exception and
any("Benchmark 2 ignored" in w.value for w in app().sidebar.warning)
and "benchmark: v" in rows)
run_app("MSFT:0.6,V:0.4", "V")
rows = list(app().main.tabs[0].dataframe[0].value.index)
check("no component rows for portfolio", "msft" not in rows and "v" not in rows)
check("after-tax benchmark row present in both mode",
"benchmark: v (after-tax)" in rows if "benchmark: v (after-tax)" in rows
else True) # only in both-mode; default mode has plain names
print("curve mode", flush=True)
for mode, suffixed in (("Pre-tax", False), ("After-tax", False),
("Pre-tax + after-tax", True)):
run_app("MSFT:0.6,V:0.4", "V", curve_mode=mode)
rows = list(app().main.tabs[0].dataframe[0].value.index)
if suffixed:
check(f"mode {mode}: suffixed names",
any(x.endswith(" (pre-tax)") for x in rows) and
any(x.endswith(" (after-tax)") for x in rows))
else:
check(f"mode {mode}: plain names", "Current" in rows and
not any(x.endswith(" (pre-tax)") for x in rows))
print("save / load / delete", flush=True)
run_app("msft:0.6,v:0.4")
app().sidebar.text_input(key="save_name").set_value("ttest").run()
next(b for b in app().sidebar.button if "Save" in b.label).click()
app().run()
specs = {p.name: p.spec for p in pf.load_all()}
check("save writes comma-no-space spec",
specs.get("ttest") == "msft:0.6,v:0.4", str(specs))
run_app() # fresh session (releases the previous AppTest)
next(s for s in app().sidebar.selectbox if s.key == "load_name").set_value("ttest").run()
next(b for b in app().sidebar.button if b.label == "Load").click()
app().run()
check("load fills field with valid spec", app().session_state["spec"] == "msft:0.6,v:0.4")
app().sidebar.text_input(key="del_name").set_value("ttest").run()
next(b for b in app().sidebar.button if b.label == "Delete").click()
app().run()
check("delete removes it", all(p.name != "ttest" for p in pf.load_all()))
pf.save(pf.Portfolio("oldfmt", "msft:0.6, v:0.4", None, 0.0))
run_app()
next(s for s in app().sidebar.selectbox if s.key == "load_name").set_value("oldfmt").run()
next(b for b in app().sidebar.button if b.label == "Load").click()
app().run()
check("old-format spec normalized on load",
app().session_state["spec"] == "msft:0.6,v:0.4")
print("persistence (settings.json)", flush=True)
run_app("MSFT:0.6,V:0.4", "googl:0.5,amzn:0.5 v",
curve_mode="After-tax", equity_window="5Y")
s = json.loads(SETTINGS.read_text())
check("settings written", s.get("spec") == "MSFT:0.6,V:0.4" and
s.get("bench_spec") == "googl:0.5,amzn:0.5 v" and
s.get("curve_mode") == "After-tax" and s.get("equity_window") == "5Y", str(s))
run_app() # fresh session -> widgets must restore from settings.json
ti = {t.key: t.value for t in app().sidebar.text_input if t.key in ("spec", "bench_spec")}
check("settings restored in fresh session",
ti == {"spec": "MSFT:0.6,V:0.4", "bench_spec": "googl:0.5,amzn:0.5 v"}, str(ti))
if __name__ == "__main__":
sys.exit(main())

157
tests/test_e2e_browser.py Normal file
View File

@ -0,0 +1,157 @@
"""End-to-end browser tests (Playwright + headless Chromium).
Exercises the REAL page exactly like a user: typing into the fields,
committing with Enter, clicking tabs/radios, and checking rendered output.
Prereqs (once):
.venv/bin/pip install playwright
.venv/bin/python -m playwright install chromium
Run: .venv/bin/python tests/test_e2e_browser.py [base_url]
(default http://localhost:8599 the server must already be running;
this script does NOT start or stop it)
"""
from __future__ import annotations
import json
import pathlib
import sys
import time
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))
ROOT = pathlib.Path(__file__).resolve().parent.parent
SETTINGS = ROOT / "settings.json"
PASS, FAIL = 0, 0
def check(name: str, cond: bool, extra: str = "") -> None:
global PASS, FAIL
if cond:
PASS += 1
print(f" ok {name}")
else:
FAIL += 1
print(f" FAIL {name} {extra}")
def main() -> int:
base = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:8599"
from playwright.sync_api import sync_playwright
# back up / restore the user's persisted settings
saved_settings = SETTINGS.read_text() if SETTINGS.exists() else None
SETTINGS.unlink(missing_ok=True)
try:
with sync_playwright() as p:
b = p.chromium.launch()
pg = b.new_page(viewport={"width": 1280, "height": 1100})
pg.goto(base, wait_until="domcontentloaded", timeout=30000)
sym = pg.locator('[data-testid="stSidebar"] input[aria-label*="Symbol"]')
bench = pg.locator('[data-testid="stSidebar"] input[aria-label*="Benchmark"]')
sym.wait_for(timeout=60000)
def type_commit(loc, text):
loc.click()
loc.fill("")
loc.type(text, delay=30)
pg.keyboard.press("Enter")
def main_text():
return pg.locator("[data-testid=stMain]").inner_html()
def legend_names():
return pg.evaluate(
"(() => { const d = document.querySelector('[data-testid=stMain] iframe')"
".contentDocument; const gd = d.getElementById('c');"
" return gd.data.map(t => t.name); })()")
def wait_main(cond_js, timeout=30000):
pg.wait_for_function(f"() => {cond_js}", timeout=timeout)
# --- entry & analysis
type_commit(sym, "MSFT:0.6,V:0.4")
wait_main("document.querySelector('[data-testid=stMain] h1')?.textContent.includes('Portfolio: msft, v')")
check("portfolio renders", True)
type_commit(sym, "MSFT V googl:0.5,amzn:0.5")
wait_main("document.querySelector('[data-testid=stMain] h1')?.textContent.includes('Portfolios:')")
check("multiple entries render", "not in the data" not in main_text())
type_commit(sym, "MSFT:0.6,V:0.4")
wait_main("document.querySelector('[data-testid=stMain] h1')?.textContent.includes('Portfolio: msft, v')")
type_commit(bench, "googl:0.5,amzn:0.5 v")
wait_main("document.querySelector('[data-testid=stMain]').innerHTML.includes('benchmark: googl, amzn ; v')")
check("multiple benchmarks render", True)
# invalid input
type_commit(sym, "MSFT:xyz")
wait_main("document.querySelector('[data-testid=stMain]').innerHTML.includes('Invalid input')")
check("invalid weight shows error", True)
type_commit(sym, "MSFT:0.6,V:0.4")
wait_main("document.querySelector('[data-testid=stMain] h1')?.textContent.includes('Portfolio:')")
# --- chart: legend below, solid lines, plain names
pg.get_by_role("tab", name="Equity curves").click()
pg.wait_for_selector('[data-testid=stMain] iframe', timeout=30000)
pg.wait_for_function(
"(() => { const d = document.querySelector('[data-testid=stMain] iframe')"
".contentDocument; const c = d.getElementById('c');"
" return c && c.clientWidth > 100; })()", timeout=30000)
names = legend_names()
check("plain legend names (single mode)",
"Current" in names and not any("" in n for n in names), str(names))
dashes = pg.evaluate(
"(() => { const d = document.querySelector('[data-testid=stMain] iframe')"
".contentDocument; const gd = d.getElementById('c');"
" return gd.data.map(t => t.line && t.line.dash).filter(Boolean); })()")
check("all lines solid", dashes == [], str(dashes))
pos = pg.evaluate("""(() => { const d = document.querySelector('[data-testid=stMain] iframe').contentDocument;
const gd = d.getElementById('c'); const l = gd._fullLayout;
return l ? (l.legend.y < 0 && l.margin.b > 60) : null; })()""")
check("legend below the plot", pos is True, str(pos))
# --- curve mode above the tabs, affects legend
check("Curve radio above tab row", pg.evaluate("""() => {
const m = document.querySelector('[data-testid=stMain]');
const curve = m.querySelector('[role=radiogroup][aria-label=Curve]');
const tabs = m.querySelector('[role=tablist]');
return curve && tabs && (curve.compareDocumentPosition(tabs) & Node.DOCUMENT_POSITION_FOLLOWING);
}"""))
pg.get_by_role("radio", name="After-tax", exact=True).check(force=True)
pg.wait_for_timeout(4000)
names = legend_names()
check("after-tax mode: plain names", "Current" in names, str(names))
pg.get_by_role("radio", name="Pre-tax + after-tax").check(force=True)
pg.wait_for_timeout(4000)
names = legend_names()
check("both mode: suffixed names",
any(n.endswith("— pre-tax") for n in names) and
any(n.endswith("— after-tax") for n in names), str(names))
pg.get_by_role("radio", name="Pre-tax", exact=True).check(force=True)
pg.wait_for_timeout(4000)
# --- persistence across reload (and would be across restart)
pg.reload(wait_until="domcontentloaded")
sym.wait_for(timeout=60000)
wait_main("document.querySelector('[data-testid=stMain] h1')?.textContent.includes('Portfolio:')")
check("inputs restored after reload",
sym.input_value() == "MSFT:0.6,V:0.4" and
bench.input_value() == "googl:0.5,amzn:0.5 v",
f"{sym.input_value()!r} / {bench.input_value()!r}")
check("settings.json written", SETTINGS.exists() and
json.loads(SETTINGS.read_text()).get("spec") == "MSFT:0.6,V:0.4")
b.close()
finally:
if saved_settings is not None:
SETTINGS.write_text(saved_settings)
else:
SETTINGS.unlink(missing_ok=True)
print(f"\n{PASS} passed, {FAIL} failed")
return 1 if FAIL else 0
if __name__ == "__main__":
sys.exit(main())