diff --git a/marimo/_convert/notebook.py b/marimo/_convert/notebook.py index f200520be97..3b20fc9a062 100644 --- a/marimo/_convert/notebook.py +++ b/marimo/_convert/notebook.py @@ -1,6 +1,8 @@ # Copyright 2026 Marimo. All rights reserved. from __future__ import annotations +from marimo._ast.cell_id import CellIdGenerator +from marimo._ast.names import SETUP_CELL_NAME from marimo._schemas.notebook import ( NotebookCell, NotebookCellConfig, @@ -12,6 +14,7 @@ CellDef, NotebookSerialization, NotebookSerializationV1, + SetupCell, ) from marimo._utils.code import hash_code from marimo._version import __version__ @@ -28,11 +31,18 @@ def convert_from_ir_to_notebook_v1( Returns: NotebookV1: The notebook v1. """ + cell_id_generator = CellIdGenerator() cells: list[NotebookCell] = [] - for data in notebook_ir.cells: + for i, data in enumerate(notebook_ir.cells): + if isinstance(data, SetupCell) or ( + i == 0 and data.name == SETUP_CELL_NAME + ): + cell_id = SETUP_CELL_NAME + else: + cell_id = cell_id_generator.create_cell_id() cells.append( NotebookCell( - id=None, + id=cell_id, code=data.code, code_hash=hash_code(data.code) if data.code else None, name=data.name, diff --git a/tests/_convert/test_notebook_cell_ids.py b/tests/_convert/test_notebook_cell_ids.py new file mode 100644 index 00000000000..d0ad954d704 --- /dev/null +++ b/tests/_convert/test_notebook_cell_ids.py @@ -0,0 +1,125 @@ +# Copyright 2026 Marimo. All rights reserved. +from __future__ import annotations + +from textwrap import dedent + +from marimo._ast.cell_id import CellIdGenerator +from marimo._ast.names import SETUP_CELL_NAME +from marimo._convert.converters import MarimoConvert + + +def _kernel_cell_ids(source: str) -> list[str]: + """Simulate the cell IDs the kernel would generate for a notebook.""" + from marimo._ast.parse import parse_notebook + from marimo._schemas.serialization import SetupCell + + ir = parse_notebook(source) + assert ir is not None + gen = CellIdGenerator() + ids = [] + for i, cell_def in enumerate(ir.cells): + if isinstance(cell_def, SetupCell) or ( + i == 0 and cell_def.name == SETUP_CELL_NAME + ): + ids.append(SETUP_CELL_NAME) + else: + ids.append(gen.create_cell_id()) + return ids + + +def test_snapshot_ids_match_kernel_ids(): + source = dedent(''' + import marimo + + __generated_with = "0.1.0" + app = marimo.App() + + @app.cell + def hello(): + x = 1 + return (x,) + + @app.cell + def world(x): + y = x + 1 + return (y,) + + if __name__ == "__main__": + app.run() + ''').strip() + + notebook = MarimoConvert.from_py(source).to_notebook_v1() + snapshot_ids = [c["id"] for c in notebook["cells"]] + kernel_ids = _kernel_cell_ids(source) + + assert snapshot_ids == kernel_ids + assert all(id is not None for id in snapshot_ids) + + +def test_snapshot_ids_match_kernel_ids_with_setup_cell(): + source = dedent(''' + import marimo + + __generated_with = "0.1.0" + app = marimo.App() + + with app.setup: + import numpy as np + + @app.cell + def hello(): + x = 1 + return (x,) + + @app.cell + def world(x): + y = x + 1 + return (y,) + + if __name__ == "__main__": + app.run() + ''').strip() + + notebook = MarimoConvert.from_py(source).to_notebook_v1() + snapshot_ids = [c["id"] for c in notebook["cells"]] + kernel_ids = _kernel_cell_ids(source) + + assert snapshot_ids == kernel_ids + assert snapshot_ids[0] == SETUP_CELL_NAME + + +def test_snapshot_ids_are_deterministic(): + source = dedent(''' + import marimo + + __generated_with = "0.1.0" + app = marimo.App() + + @app.cell + def _(): + x = 1 + return (x,) + + @app.cell + def _(): + y = 2 + return (y,) + + @app.cell + def _(): + z = 3 + return (z,) + + if __name__ == "__main__": + app.run() + ''').strip() + + ids_1 = [ + c["id"] for c in MarimoConvert.from_py(source).to_notebook_v1()["cells"] + ] + ids_2 = [ + c["id"] for c in MarimoConvert.from_py(source).to_notebook_v1()["cells"] + ] + + assert ids_1 == ids_2 + assert len(set(ids_1)) == 3 # all unique