From a3f06fc7edc60a976ad221a6475536fc606e8d94 Mon Sep 17 00:00:00 2001 From: dylan Date: Thu, 16 Oct 2025 15:58:23 -0700 Subject: [PATCH 1/4] fix: allow cache blocks to be invoked in script mode --- marimo/_ast/load.py | 18 +++++++++++++++++- marimo/_save/save.py | 25 +++++++++++++++++++++++-- tests/_save/external_decorators/app.py | 11 ++++++++++- tests/_save/test_external_decorators.py | 4 ++++ 4 files changed, 54 insertions(+), 4 deletions(-) diff --git a/marimo/_ast/load.py b/marimo/_ast/load.py index c3e6f77452e..e40188e345e 100644 --- a/marimo/_ast/load.py +++ b/marimo/_ast/load.py @@ -14,7 +14,11 @@ is_non_marimo_python_script, parse_notebook, ) -from marimo._schemas.serialization import NotebookSerialization, UnparsableCell +from marimo._schemas.serialization import ( + CellDef, + NotebookSerialization, + UnparsableCell, +) LOGGER = _loggers.marimo_logger() @@ -89,6 +93,18 @@ def _static_load(filepath: Path) -> Optional[App]: return load_notebook_ir(notebook, filepath=str(filepath)) +def find_cell(filename, lineno) -> CellDef: + load_result = get_notebook_status(filename) + if load_result.notebook is None: + raise OSError("Could not resolve notebook.") + previous = None + for cell in load_result.notebook.cells: + if cell.lineno > lineno: + break + previous = cell + return previous + + def load_notebook_ir( notebook: NotebookSerialization, filepath: Optional[str] = None ) -> App: diff --git a/marimo/_save/save.py b/marimo/_save/save.py index bd6e9795e8f..bff421435bb 100644 --- a/marimo/_save/save.py +++ b/marimo/_save/save.py @@ -27,6 +27,7 @@ ) from marimo._ast.cell_id import is_external_cell_id +from marimo._ast.load import find_cell from marimo._ast.transformers import ( ARG_PREFIX, CacheExtractWithBlock, @@ -596,7 +597,7 @@ def trace(self, with_frame: FrameType) -> None: # causing this function to terminate before reaching this block. self._frame = with_frame for i, frame in enumerate(stack[::-1]): - _filename, lineno, function_name, _code = frame + filename, lineno, function_name, _code = frame if function_name == "": ctx = get_context() if ctx.execution_context is None: @@ -608,10 +609,30 @@ def trace(self, with_frame: FrameType) -> None: ) graph = ctx.graph cell_id = ctx.cell_id or ctx.execution_context.cell_id + + # We are calling from script mode, so our line number is + # absolute. + if "__marimo__" not in filename: + cell = find_cell(filename, lineno) + if cell is None: + raise CacheException( + "Could not resolve cell for cache." + f"{UNEXPECTED_FAILURE_BOILERPLATE}" + ) + lineno -= cell.lineno + code = cell.code + elif cell_id in graph.cells: + code = graph.cells[cell_id].code + else: + raise CacheException( + "Could not resolve cell for cache." + f"{UNEXPECTED_FAILURE_BOILERPLATE}" + ) + pre_module, save_module = CacheExtractWithBlock( lineno - 1 ).visit( - ast.parse(graph.cells[cell_id].code).body # type: ignore[arg-type] + ast.parse(code).body # type: ignore[arg-type] ) self._cache = cache_attempt_from_hash( diff --git a/tests/_save/external_decorators/app.py b/tests/_save/external_decorators/app.py index 4d999919a8b..7f07522fb0f 100644 --- a/tests/_save/external_decorators/app.py +++ b/tests/_save/external_decorators/app.py @@ -8,7 +8,7 @@ @app.cell -def _(): +def decorator_wrap(): @mo.cache def cache(x): return x + 1 @@ -17,5 +17,14 @@ def cache(x): return (bar, cache) +@app.cell +def block_wrap(mo): + with mo.cache("random"): + x = [] + + a = "need a final line to trigger invalid block capture" + return (x,) + + if __name__ == "__main__": app.run() diff --git a/tests/_save/test_external_decorators.py b/tests/_save/test_external_decorators.py index 0f180500673..6fe62646986 100644 --- a/tests/_save/test_external_decorators.py +++ b/tests/_save/test_external_decorators.py @@ -282,6 +282,10 @@ def _(): _, defs = ex_app.run() assert defs["bar"] == 2 assert defs["cache"](1) == 2 + assert len(defs["x"]) == 0 + defs["x"].append(1) + _, defs = ex_app.run() + assert len(defs["x"]) == 1 return @staticmethod From e14be539d382c900d2a1dbf5587318d0e2fb0d95 Mon Sep 17 00:00:00 2001 From: dylan Date: Thu, 16 Oct 2025 16:28:29 -0700 Subject: [PATCH 2/4] fix: typing --- marimo/_ast/load.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/marimo/_ast/load.py b/marimo/_ast/load.py index e40188e345e..e87231a1f74 100644 --- a/marimo/_ast/load.py +++ b/marimo/_ast/load.py @@ -93,7 +93,7 @@ def _static_load(filepath: Path) -> Optional[App]: return load_notebook_ir(notebook, filepath=str(filepath)) -def find_cell(filename, lineno) -> CellDef: +def find_cell(filename, lineno) -> CellDef | None: load_result = get_notebook_status(filename) if load_result.notebook is None: raise OSError("Could not resolve notebook.") From 850960f5948be6c93f820dc4bdbf6109eb60ab86 Mon Sep 17 00:00:00 2001 From: dylan Date: Fri, 17 Oct 2025 11:54:35 -0700 Subject: [PATCH 3/4] typing: missing args --- marimo/_ast/load.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/marimo/_ast/load.py b/marimo/_ast/load.py index e87231a1f74..aadf86c9c7e 100644 --- a/marimo/_ast/load.py +++ b/marimo/_ast/load.py @@ -93,7 +93,12 @@ def _static_load(filepath: Path) -> Optional[App]: return load_notebook_ir(notebook, filepath=str(filepath)) -def find_cell(filename, lineno) -> CellDef | None: +def find_cell(filename: str, lineno: int) -> CellDef | None: + """Find the cell at the given line number in the notebook. + Args: + filename: Path to a marimo notebook file (.py or .md) + lineno: Line number to search for + """ load_result = get_notebook_status(filename) if load_result.notebook is None: raise OSError("Could not resolve notebook.") From 995581947b5a6300c2d9335187082bff2164244f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 17 Oct 2025 18:55:39 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- marimo/_ast/load.py | 1 + 1 file changed, 1 insertion(+) diff --git a/marimo/_ast/load.py b/marimo/_ast/load.py index aadf86c9c7e..b207adc987b 100644 --- a/marimo/_ast/load.py +++ b/marimo/_ast/load.py @@ -95,6 +95,7 @@ def _static_load(filepath: Path) -> Optional[App]: def find_cell(filename: str, lineno: int) -> CellDef | None: """Find the cell at the given line number in the notebook. + Args: filename: Path to a marimo notebook file (.py or .md) lineno: Line number to search for