Coverage for src/jointview/data.py: 100%
43 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-17 06:44 +0000
1"""Getting a DataFrame into the app.
3The app is handed a path on the command line; when there is none it falls back to
4a generated frame so that ``marimo run app.py`` works out of the box.
5"""
7from __future__ import annotations
9from collections.abc import Callable
10from datetime import date, timedelta
11from pathlib import Path
13import numpy as np
14import polars as pl
16# Every entry has to hand back a frame whose dates are dates. The binary formats carry
17# a schema and do it for free; the text ones cannot say that a column of "2024-01-01"
18# was ever anything but text, so each says here how it recovers them. It matters beyond
19# the x-axis: `stats` reads the annualisation factor off the spacing of the period
20# column, and a date column arriving as text costs a weekly series its real one.
21READERS: dict[str, Callable[[Path], pl.DataFrame]] = {
22 ".arrow": pl.read_ipc,
23 ".csv": lambda p: pl.read_csv(p, try_parse_dates=True),
24 ".feather": pl.read_ipc,
25 ".ipc": pl.read_ipc,
26 # No try_parse_dates on the JSON readers — polars offers the hook for CSV only — so
27 # the same promotion is done on the way out instead.
28 ".json": lambda p: _with_dates(pl.read_json(p)),
29 ".ndjson": lambda p: _with_dates(pl.read_ndjson(p)),
30 ".parquet": pl.read_parquet,
31 ".tsv": lambda p: pl.read_csv(p, separator="\t", try_parse_dates=True),
32}
34# name: starting level, sensitivity to the common market move, daily drift of its
35# own, and the size of the wobble nobody else shares.
36FUNDS: dict[str, tuple[float, float, float, float]] = {
37 "world_equity": (100.0, 1.00, 0.00000, 0.0030),
38 "tech_fund": (48.5, 1.35, 0.00030, 0.0090),
39 "value_fund": (212.0, 0.85, 0.00005, 0.0060),
40 "balanced": (1_450.0, 0.45, 0.00010, 0.0030),
41 "bond_fund": (98.0, 0.10, 0.00004, 0.0020),
42 "cash": (1.0, 0.00, 0.00008, 0.00002),
43}
46def load_frame(path: str | Path | None) -> pl.DataFrame:
47 """Read a DataFrame from ``path``, or build the demo frame when it is None.
49 The suffix picks the reader, so no path means the generated frame rather than
50 an error — which is what makes ``marimo run app.py`` work with no arguments:
52 >>> load_frame(None).columns[0]
53 'date'
55 Whatever the format, a column of dates arrives as dates: the text formats have no
56 way to record that one ever was, so they are parsed back on the way in. A column
57 that does not read as dates all the way down is left as the text it is.
59 A path that is not there is reported before a reader is chosen, so the message
60 names the file rather than complaining about its extension:
62 >>> load_frame("nowhere.parquet")
63 Traceback (most recent call last):
64 ...
65 FileNotFoundError: no such file: nowhere.parquet
67 A directory is named as one. This is also where an empty ``Path`` lands, since
68 ``Path("")`` is ``Path(".")`` — the two are the same object by the time anything
69 here sees them, so the current directory is what actually arrived:
71 >>> load_frame(".")
72 Traceback (most recent call last):
73 ...
74 IsADirectoryError: not a file: .
75 """
76 # Only the empty *string* is the "no path" sentinel: it is what marimo hands over
77 # for a flag that was not passed. A Path cannot carry it — see the docstring.
78 if path is None or path == "":
79 return demo_frame()
81 file = Path(path).expanduser()
82 if not file.exists():
83 raise FileNotFoundError(f"no such file: {file}") # noqa: TRY003
85 # Before the suffix lookup, which would otherwise reject a directory for having
86 # the wrong extension — and an empty one for having no name to quote at all.
87 if file.is_dir():
88 raise IsADirectoryError(f"not a file: {file}") # noqa: TRY003
90 reader = READERS.get(file.suffix.lower())
91 if reader is None:
92 supported = ", ".join(sorted(READERS))
93 raise ValueError(f"cannot read {file.suffix or file.name!r}; supported: {supported}") # noqa: TRY003
95 return reader(file)
98def demo_frame(rows: int = 1_500, seed: int = 42) -> pl.DataFrame:
99 """Daily NAVs for a handful of made-up funds, on deliberately different scales.
101 They share a market factor, so the lines rhyme without being copies, and they
102 start anywhere from 1 to 1,450 — which is exactly the case that needs indexing
103 before two of them can be read on one axis.
105 The dates come first, so the frame is ready for :func:`jointview.plot.line_chart`
106 as it stands, and ``seed`` makes it the same frame every time:
108 >>> frame = demo_frame(rows=10)
109 >>> frame.columns[:2]
110 ['date', 'world_equity']
111 >>> frame.height
112 10
113 """
114 rng = np.random.default_rng(seed)
115 market = rng.normal(0.0004, 0.011, rows)
117 navs = {
118 name: start * np.cumprod(1.0 + drift + beta * market + rng.normal(0.0, wobble, rows))
119 for name, (start, beta, drift, wobble) in FUNDS.items()
120 }
121 return pl.DataFrame({"date": _business_days(date(2020, 1, 1), rows), **navs})
124def _business_days(start: date, rows: int) -> pl.Series:
125 """``rows`` weekdays from ``start``, so 252 periods really are about a year."""
126 days = pl.date_range(start, start + timedelta(days=2 * rows), "1d", eager=True)
127 return days.filter(days.dt.weekday() <= 5).head(rows).alias("date")
130def _with_dates(frame: pl.DataFrame) -> pl.DataFrame:
131 """``frame`` with every text column that reads cleanly as dates promoted to dates.
133 What :func:`pl.read_csv`'s ``try_parse_dates`` does during the parse, done after it
134 — the JSON readers take no such flag, and ``schema_overrides`` would need the column
135 named in advance, which is exactly what is not known here.
136 """
137 promoted = []
138 for name, dtype in frame.schema.items():
139 if dtype == pl.String:
140 parsed = _as_dates(frame[name])
141 if parsed is not None:
142 promoted.append(parsed.alias(name))
143 return frame.with_columns(promoted)
146def _as_dates(column: pl.Series) -> pl.Series | None:
147 """``column`` read as dates, or None where it is text that merely looks like some.
149 The bar is the whole column: a fund name, a code that happens to be eight digits,
150 or a date column with one bad row all stay as they arrived. Promoting on a partial
151 match would turn the rows that failed into nulls, and a null in the period column
152 is a row silently dropped from both the plot and the summary beside it.
153 """
154 try:
155 parsed = column.str.to_date(strict=False)
156 except pl.exceptions.ComputeError:
157 # Raised when no format fits *any* value — ordinary text. `strict=False` governs
158 # the values that fail once a format is chosen, not the choosing of it.
159 return None
160 # A value that did not parse comes back null, so an unchanged null count is the test
161 # for "every row was a date". The second clause is what stops a column of nothing
162 # becoming a column of no dates: polars hands an all-null column over as `Null`
163 # rather than `String` today, so nothing reaches here to need it, but a dateless
164 # date column would be a poor thing to acquire on the strength of that.
165 return parsed if parsed.null_count() == column.null_count() < parsed.len() else None