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 

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

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

29 

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

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

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

33 """ 

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

35 

36 

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

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

39 path = Path(notebook_file).parent / "public" / "Prices_hashed.csv" 

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

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

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

43 return interpolate(dframe)