Coverage for book/marimo/notebooks/preamble.py: 100%

1"""Shared data loading for marimo experiment notebooks. 

2 

3This module also hosts :func:`load_notebook`, the single place that executes a 

4sibling notebook (or ``optimize.py``) via :func:`runpy.run_path` and returns its 

5namespace. The experiment notebooks are not an importable package, so both 

6``optimize.py`` and the test suite need to read symbols (the signal ``f``, the 

7``build_exp*`` builders, …) out of a freshly executed notebook namespace; 

8centralizing that here keeps the ``runpy`` call in one place. 

9""" 

10 

11import runpy 

12from pathlib import Path 

13from typing import Any 

14 

15import plotly.io as pio 

16import polars as pl 

17from jquantstats import interpolate 

18 

19pio.renderers.default = "plotly_mimetype" 

20 

21date_col = "date" 

22 

23#: Directory holding the marimo notebooks (this file's own directory). 

24NOTEBOOK_DIR = Path(__file__).resolve().parent 

25 

26#: Price file every notebook reads, from the ``public/`` directory beside it. 

27PRICES_CSV = "Prices_hashed.csv" 

28 

29 

30def load_notebook(name: str) -> dict[str, Any]: 

31 """Execute sibling notebook ``name`` (e.g. ``"Experiment1.py"``) and return its namespace. 

32 

33 The returned dict maps top-level names defined by the notebook to their 

34 values, so callers can pull out the signal function with 

35 ``load_notebook("Experiment1.py")["f"]``. 

36 

37 Executing this module itself is the cheapest demonstration of that contract — 

38 the shared helpers come back as ordinary entries in the namespace: 

39 

40 >>> namespace = load_notebook("preamble.py") 

41 >>> sorted(name for name in namespace if name in {"date_col", "load_prices"}) 

42 ['date_col', 'load_prices'] 

43 """ 

44 return runpy.run_path(str(NOTEBOOK_DIR / name)) 

45 

46 

47def load_prices(notebook_file: str) -> pl.DataFrame: 

48 """Load and preprocess prices from the standard CSV file. 

49 

50 ``notebook_file`` is the *caller's* own path — the notebooks pass ``__file__`` — 

51 and the CSV is read from its ``public/`` sibling directory. The frame comes back 

52 with the date column first as nanosecond datetimes, every remaining column (one 

53 per asset) cast to ``Float64``, and gaps interpolated: 

54 

55 >>> prices = load_prices(str(NOTEBOOK_DIR / "preamble.py")) 

56 >>> prices.columns[0] 

57 'date' 

58 >>> prices[date_col].dtype 

59 Datetime(time_unit='ns', time_zone=None) 

60 >>> set(prices.drop(date_col).dtypes) == {pl.Float64} 

61 True 

62 

63 An absent CSV is checked for here rather than left to ``pl.read_csv``, so the 

64 error names the file that was looked for — the file ships with the repository, 

65 so its absence is a setup problem with an obvious remedy: 

66 

67 >>> try: 

68 ... load_prices(str(NOTEBOOK_DIR / "elsewhere" / "preamble.py")) 

69 ... except FileNotFoundError as error: 

70 ... PRICES_CSV in str(error) 

71 True 

72 """ 

73 path = Path(notebook_file).parent / "public" / PRICES_CSV 

74 if not path.is_file(): 

75 msg = f"Price data not found: {path} — it ships with the repository; check the checkout is complete." 

76 raise FileNotFoundError(msg) 

77 dframe = pl.read_csv(str(path), try_parse_dates=True) 

78 dframe = dframe.with_columns(pl.col(date_col).cast(pl.Datetime("ns"))) 

79 dframe = dframe.with_columns([pl.col(col).cast(pl.Float64) for col in dframe.columns if col != date_col]) 

80 return interpolate(dframe)