Merge pull request #504 from openclimatefix/claude/github-issue-500-review-855247

Catch marimo notebooks whose cells reference unbound names

Author: JackKelly

.claude/skills/marimo-notebooks/SKILL.md

@@ -1,13 +1,15 @@
 ---
 name: marimo-notebooks
 description: >-
-  Three authoring rules for this repo's Marimo notebooks (`packages/dashboard/*.py`,
-  `packages/notebooks/*.py`), each reversing a normal Python habit: a leading underscore makes a
-  name cell-local, so cross-cell helpers must be public; every import belongs in `with app.setup:`
-  or ruff stops seeing it; and `ruff check --fix` must never be run over a notebook, because an
-  autofix can insert an import outside `app.setup` and break it while reporting success. Load
-  before creating or editing a Marimo notebook, or when one fails with a `NameError` on a helper
-  or an import that looks present.
+  Authoring rules for this repo's Marimo notebooks (`packages/dashboard/*.py`,
+  `packages/notebooks/*.py`), most of them reversing a normal Python habit: a leading underscore
+  makes a name cell-local, so cross-cell helpers must be public; every import belongs in `with
+  app.setup:` or ruff stops seeing it; `ruff check --fix` must never be run over a notebook,
+  because an autofix can insert an import outside `app.setup` and break it while reporting
+  success; and a helper belongs in the `@app.function` form marimo itself writes. Load before
+  creating or editing a Marimo notebook, when one fails with a `NameError` on a helper or an
+  import that looks present, or when the `check-marimo-notebooks` hook reports a cell referencing
+  a name no cell defines.
 ---
 
 # Authoring Marimo notebooks
@@ -47,3 +49,21 @@ does it, and ruff has no per-file fixability setting to prevent it (`unfixable`
 The pre-commit hook is split so notebooks are checked but never auto-fixed; a bare `uv run ruff
 check . --fix` typed by hand is *not* covered, so after running one, check `git diff` for an
 import that landed above `import marimo` and move it into `app.setup`.
+
+`marimo check --fix` does not rescue a notebook in that state either: it deletes the module-level
+import and rewrites the cell that used the name as `def _(name)`, which leaves the name as a cell
+input nothing defines — broken in a second way.
+
+Both shapes are caught by `scripts/check_marimo_notebooks.py`, which runs as a pre-commit hook over
+changed notebooks and over every notebook from `tests/test_marimo_notebooks.py`. So a mistake here
+fails the commit or CI rather than surviving to whoever next opens the notebook. What it catches
+and what it cannot:
+<https://openclimatefix.github.io/nged-substation-forecast/architecture/testing/#marimo-notebooks-bind-every-name-their-cells-reference>
+
+## Let `marimo check --fix` settle a notebook's shape before committing
+
+A helper hand-written inside an `@app.cell` gets rewritten to a top-level `@app.function` the next
+time marimo saves the notebook, so committing the hand-written form buys a large diff for no
+change. Write helpers — and any `test_*` function exercising them — in the `@app.function` form,
+or run `marimo check --fix` before committing. `packages/notebooks/plot_missing_NWP_data.py` is
+the worked example.

.pre-commit-config.yaml

@@ -44,6 +44,16 @@ repos:
         language: system
         types: [python]
         files: *marimo_notebooks
+      # The hooks above never auto-fix a notebook, but a hand-typed `uv run ruff check . --fix`
+      # still can, and `marimo check --fix` leaves the same breakage behind by another route. So
+      # every commit that touches a notebook also checks its cells bind every name they reference
+      # — see the script's docstring. `tests/test_marimo_notebooks.py` runs the same check in CI.
+      - id: check-marimo-notebooks
+        name: marimo notebooks (no unbound names)
+        entry: uv run python scripts/check_marimo_notebooks.py
+        language: system
+        types: [python]
+        files: *marimo_notebooks
       - id: ruff-format
         name: ruff format
         entry: uv run ruff format --force-exclude

docs/architecture/testing.md

@@ -239,6 +239,42 @@ present):
 the only layer that can catch *future upstream drift* — a change in Dynamical.org's own conventions
 that the committed slice, frozen at capture time, cannot.
 
+## Marimo notebooks bind every name their cells reference
+
+Marimo rebuilds a notebook from its `with app.setup:` block plus its `@app.cell` functions, and
+never runs the module-level statements in between. A name bound at module level is therefore
+invisible to every cell: the notebook raises `NameError` the next time it is opened, while ruff, ty
+and pytest all pass, because the file they were handed is valid Python. Two tools produce exactly
+that shape from a working notebook — `ruff check --fix`, which writes an import an autofix needs
+into the top-level import block, and `marimo check --fix`, which deletes such an import and
+rewrites the cell that used the name as `def _(name)`, leaving a cell input nothing defines.
+
+`scripts/check_marimo_notebooks.py` reads each cell's `refs` and `defs` and reports any name a cell
+references that no cell binds. It runs as a pre-commit hook over changed notebooks, and
+`tests/test_marimo_notebooks.py` runs it over every notebook in `packages/notebooks/` and
+`packages/dashboard/`. Three properties are worth knowing:
+
+- **It is static.** Nothing executes, so the check needs none of the notebooks' runtime
+  dependencies — only marimo itself, which the root environment has via the `dashboard` dev
+  dependency. It cannot catch a notebook that binds every name and still fails inside a Polars or
+  Altair call; executing the notebooks is not an option, because they read real Delta tables and
+  S3.
+- **It rides on private marimo API.** `Cell.refs` and `Cell.defs` are documented, but loading a
+  notebook without running it is not. So the checker raises rather than reporting "no findings"
+  whenever a file does not parse into at least one cell, and the tests keep a positive control —
+  a deliberately broken notebook, held as a string so ruff never sees it — that fails if a marimo
+  release stops the check detecting a real breakage.
+- **Every `.py` file directly inside those two directories must be a notebook**, and a file that
+  is not one is a finding. The ruff pre-commit hooks share that assumption: they use it to decide
+  which files must never be auto-fixed.
+
+Testing what a notebook's cells actually *do* is a separate job, and
+`packages/notebooks/plot_missing_NWP_data.py` is the worked example. Its chart-building helper is
+an `@app.function` — marimo's form for a top-level reusable function — so an ordinary `test_*`
+function in the same notebook can exercise it on a synthetic frame. Naming the notebook in
+`python_files` is what makes a plain `uv run pytest` collect it. The authoring rules for writing
+one are in the `marimo-notebooks` skill.
+
 ## Assertion style for Patito frames
 
 Build a frame, attach the model, cast, and validate for the happy path:

packages/notebooks/plot_missing_NWP_data.py

@@ -1,6 +1,6 @@
 import marimo
 
-__generated_with = "0.23.6"
+__generated_with = "0.23.16"
 app = marimo.App()
 
 with app.setup:
@@ -9,13 +9,14 @@
     import altair as alt
     import plotting.ocf_theme  # noqa: F401 — registers OCF Altair theme as side effect
     import polars as pl
+    from contracts.weather_schemas import Nwp
     from plotting.ocf_theme import GRID, ORANGE_RED
 
 
 @app.cell
 def _():
-    # TODO: Load polars DataFrame from local Delta Table of NWP data.
-    return
+    df = Nwp.scan_delta().drop("nwp_model_id")
+    return (df,)
 
 
 @app.cell
@@ -38,71 +39,83 @@ def _():
     return (nwp_vars,)
 
 
-@app.cell
-def _(df, nwp_vars):
-    def plot_null_distribution(df, target_init_time, target_h3_index, nwp_vars):
-        # 1. Filter down to the specific init_time and h3_index
-        filtered = df.filter(
-            (pl.col("init_time") == target_init_time) & (pl.col("h3_index") == target_h3_index)
-        )
+@app.function
+def plot_null_distribution(df, target_init_time, target_h3_index, nwp_vars):
+    """Chart where `nwp_vars` is missing, one row per ensemble member.
 
-        # 2. Unpivot (melt) the data so we can plot all variables on a single Y-axis
-        melted = filtered.unpivot(
-            on=nwp_vars,
-            index=["valid_time", "ensemble_member"],
-            variable_name="variable",
-            value_name="value",
-        )
+    Args:
+        df: Lazy scan of the NWP Delta table.
+        target_init_time: The single NWP run to plot.
+        target_h3_index: The single H3 cell to plot.
+        nwp_vars: The NWP variable name (or list of names) to plot.
+    """
+    # 1. Filter down to the specific init_time and h3_index
+    filtered = df.filter(
+        (pl.col("init_time") == target_init_time) & (pl.col("h3_index") == target_h3_index)
+    )
 
-        # 3. Create a boolean flag for missing data and a combined label for the Y-axis
-        plot_df = melted.with_columns(
-            # Check for both database Nulls and float NaNs
-            is_missing=pl.col("value").is_null() | pl.col("value").is_nan(),
-            # Create a string like "temperature_2m (Member 0)":
-            # row_label=(
-            #     pl.col("variable") + " (Member " + pl.col("ensemble_member").cast(pl.Utf8) + ")"
-            # ),
-            row_label=pl.col("ensemble_member"),
-        )
+    # 2. Unpivot (melt) the data so we can plot all variables on a single Y-axis
+    melted = filtered.unpivot(
+        on=nwp_vars,
+        index=["valid_time", "ensemble_member"],
+        variable_name="variable",
+        value_name="value",
+    )
 
-        # 5. Build the Altair Chart
-        base = alt.Chart(plot_df).encode(
-            y=alt.Y("row_label:N", title="Ensemble Member", sort="ascending")
-        )
+    # 3. Create a boolean flag for missing data and a combined label for the Y-axis. Altair
+    # only accepts materialised data, so collect once the filter has cut the scan down to a
+    # single NWP run and a single H3 cell.
+    plot_df = melted.with_columns(
+        # Check for both database Nulls and float NaNs
+        is_missing=pl.col("value").is_null() | pl.col("value").is_nan(),
+        # Create a string like "temperature_2m (Member 0)":
+        # row_label=(
+        #     pl.col("variable") + " (Member " + pl.col("ensemble_member").cast(pl.Utf8) + ")"
+        # ),
+        row_label=pl.col("ensemble_member"),
+    ).collect()
+
+    # 5. Build the Altair Chart
+    base = alt.Chart(plot_df).encode(
+        y=alt.Y("row_label:N", title="Ensemble Member", sort="ascending")
+    )
 
-        # Layer 1: A light gray background line showing the full time series extent
-        background_line = base.mark_line(color=GRID, strokeWidth=1)
-        background_lines = background_line.encode(  # ty: ignore[unresolved-attribute]
-            x=alt.X("valid_time:T", title="Valid Time"),
-            detail="row_label:N",  # Ensures lines don't connect across different rows
-        )
+    # Layer 1: A light gray background line showing the full time series extent
+    background_line = base.mark_line(color=GRID, strokeWidth=1)
+    background_lines = background_line.encode(  # ty: ignore[unresolved-attribute]
+        x=alt.X("valid_time:T", title="Valid Time"),
+        detail="row_label:N",  # Ensures lines don't connect across different rows
+    )
 
-        # Layer 2: Red ticks superimposed exactly where the data is missing
-        missing_marks = (
-            base.transform_filter(alt.datum.is_missing)
-            .mark_tick(
-                color=ORANGE_RED,
-                thickness=3,  # Make the red mark stand out
-                size=12,  # Height of the tick mark
-            )
-            .encode(x="valid_time:T")  # ty: ignore[unresolved-attribute]  # astral-sh/ty#2520
+    # Layer 2: Red ticks superimposed exactly where the data is missing
+    missing_marks = (
+        base.transform_filter(alt.datum.is_missing)
+        .mark_tick(
+            color=ORANGE_RED,
+            thickness=3,  # Make the red mark stand out
+            size=12,  # Height of the tick mark
         )
+        .encode(x="valid_time:T")  # ty: ignore[unresolved-attribute]  # astral-sh/ty#2520
+    )
 
-        # Combine the layers and configure the chart size
-        return (
-            (background_lines + missing_marks)
-            .properties(
-                title=(
-                    f"Missing NWP Data | init_time: {target_init_time.strftime('%Y-%m-%d')}"
-                    f" | {nwp_vars} | H3: {target_h3_index}"
-                ),
-                width=800,
-                # Dynamically scales chart height based on the number of rows
-                height=alt.Step(10),
-            )
-            .configure_axis(labelFontSize=11, titleFontSize=13)
+    # Combine the layers and configure the chart size
+    return (
+        (background_lines + missing_marks)
+        .properties(
+            title=(
+                f"Missing NWP Data | init_time: {target_init_time.strftime('%Y-%m-%d')}"
+                f" | {nwp_vars} | H3: {target_h3_index}"
+            ),
+            width=800,
+            # Dynamically scales chart height based on the number of rows
+            height=alt.Step(10),
         )
+        .configure_axis(labelFontSize=11, titleFontSize=13)
+    )
+
 
+@app.cell
+def _(df, nwp_vars):
     chart = plot_null_distribution(
         df,
         target_init_time=datetime(2026, 5, 1, tzinfo=UTC),
@@ -113,5 +126,37 @@ def plot_null_distribution(df, target_init_time, target_h3_index, nwp_vars):
     return
 
 
+@app.function
+def test_plot_null_distribution_flags_nulls_and_nans():
+    init_time = datetime(2026, 5, 1, tzinfo=UTC)
+    other_init_time = datetime(2026, 5, 2, tzinfo=UTC)
+    h3_index = 599148110664433663
+    other_h3_index = 599148110664433662
+    # The last two readings are missing too, but belong to another NWP run and another H3 cell, so
+    # the filter must drop them rather than plot them.
+    readings = pl.LazyFrame(
+        {
+            "init_time": [init_time] * 4 + [other_init_time, init_time],
+            "valid_time": [datetime(2026, 5, 1, hour=h, tzinfo=UTC) for h in range(6)],
+            "ensemble_member": [0, 0, 0, 0, 0, 0],
+            "h3_index": [h3_index] * 5 + [other_h3_index],
+            "temperature_2m": pl.Series(
+                [1.0, None, float("nan"), 4.0, None, None], dtype=pl.Float32
+            ),
+        }
+    )
+
+    chart = plot_null_distribution(
+        readings,
+        target_init_time=init_time,
+        target_h3_index=h3_index,
+        nwp_vars="temperature_2m",
+    )
+
+    # Altair inlines the chart's data as its single named dataset.
+    (rows,) = chart.to_dict()["datasets"].values()
+    assert [row["is_missing"] for row in rows] == [False, True, True, False]
+
+
 if __name__ == "__main__":
     app.run()

pyproject.toml

@@ -299,6 +299,12 @@ docstring-code-line-length = 100
 # with `uv run pytest --run-network` — see docs/architecture/testing.md. The gate lives in a hook,
 # not in an `addopts` `-m "not network"`, because any caller-supplied `-m` would silently override it.
 addopts = "--import-mode=importlib"
+# The first two patterns are pytest's default, restated because naming any pattern replaces the
+# default list. The notebook is named as well because a marimo notebook is an ordinary Python file
+# that pytest can collect, and this one holds a `test_*` function exercising its chart-building
+# helper. See docs/architecture/testing.md. Patterns containing a `/` match against the path
+# rather than the basename, so this reaches exactly one file.
+python_files = ["test_*.py", "*_test.py", "packages/notebooks/plot_missing_NWP_data.py"]
 # Put the root `tests/` dir on sys.path so its integration tests can import shared, non-fixture
 # test data (e.g. `_nwp_test_data`) by bare name — importlib mode does not add each test file's
 # own directory to sys.path the way prepend mode does.