From b68d10bcd763934b5c92bfed4ed8f44d0155b38f Mon Sep 17 00:00:00 2001 From: dmadisetti Date: Wed, 27 May 2026 10:10:05 -0700 Subject: [PATCH 1/2] test: add a basic test for testign private variables --- tests/_save/test_cache.py | 56 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/tests/_save/test_cache.py b/tests/_save/test_cache.py index 8c9db07faa8..d80272815c0 100644 --- a/tests/_save/test_cache.py +++ b/tests/_save/test_cache.py @@ -1111,6 +1111,62 @@ def foo(): assert k.stdout.messages.count("ran") == 1 assert k.globals["foo"].hits == 3 + async def test_lru_cache_underscore_alias( + self, k: Kernel, exec_req: ExecReqProvider + ) -> None: + # Regression: `import marimo as _private` (underscore alias) used to + # blow up cache hashing because mangled scope keys (`_cell___private`) + # don't match the unmangled-ref lookup used during the cache attempt. + await k.run( + [ + exec_req.get( + """ + import marimo as _private + + @_private.lru_cache(maxsize=128) + def f(x): + return x * 2 + + a = f(1) + b = f(1) + c = f(2) + """ + ), + ] + ) + + assert not k.stderr.messages, k.stderr + assert k.globals["a"] == 2 + assert k.globals["b"] == 2 + assert k.globals["c"] == 4 + # Second `f(1)` should hit the cache. + assert k.globals["f"].hits == 1 + + async def test_cache_underscore_alias( + self, k: Kernel, exec_req: ExecReqProvider + ) -> None: + await k.run( + [ + exec_req.get( + """ + import marimo as _private + + @_private.cache + def f(x): + return x * 3 + + a = f(1) + b = f(1) + """ + ), + ] + ) + + assert not k.stderr.messages, k.stderr + assert k.globals["a"] == 3 + assert k.globals["b"] == 3 + assert k.globals["f"].hits == 1 + async def test_persistent_cache( self, k: Kernel, exec_req: ExecReqProvider ) -> None: From ef3ace827d85afd88f3cf2f364032b8ca6a1a0fb Mon Sep 17 00:00:00 2001 From: dmadisetti Date: Wed, 27 May 2026 15:44:24 -0700 Subject: [PATCH 2/2] fix: handle private imports in general --- marimo/_ast/compiler.py | 13 ++++++++- marimo/_ast/visitor.py | 29 ++++++++++++-------- marimo/_save/hash.py | 17 +++++++++--- tests/_ast/test_visitor.py | 31 +++++++++++++++++++++ tests/_runtime/test_runtime.py | 49 +++++++++++++++++++++++++++++++--- 5 files changed, 121 insertions(+), 18 deletions(-) diff --git a/marimo/_ast/compiler.py b/marimo/_ast/compiler.py index ef4028be4cd..e45dd02e9c1 100644 --- a/marimo/_ast/compiler.py +++ b/marimo/_ast/compiler.py @@ -356,7 +356,18 @@ def compile_cell( expr, filename, mode="eval", dont_inherit=True, flags=flags ) - nonlocals = {name for name in v.defs if not is_local(name)} + # Imports are exempt from the underscore "cell-local" rule: their + # name comes from the package (the user can't control it), and we + # need them in the graph so consumers in other cells can resolve. + nonlocals = { + name + for name in v.defs + if not is_local(name) + or ( + name in v.variable_data + and any(d.kind == "import" for d in v.variable_data[name]) + ) + } temporaries = v.defs - nonlocals variable_data = { name: v.variable_data[name] diff --git a/marimo/_ast/visitor.py b/marimo/_ast/visitor.py index 6b23ec78af9..126aed13a6c 100644 --- a/marimo/_ast/visitor.py +++ b/marimo/_ast/visitor.py @@ -270,7 +270,9 @@ def _get_alias_name( ) -> str: """Get the string name of an imported alias. - Mangles the "as" name if it's a local variable. + Imports are never mangled. The user can't control upstream + package names, and `import x as _y` already gets cell-local + semantics via `is_local` -> `temporaries` in `compile_cell`. NB: We disallow `import *` because Python only allows star imports at module-level, but we store cells as functions. @@ -280,9 +282,6 @@ def _get_alias_name( # import [a.b.c] - we define a # from foo import [a] - we define a # from foo import [*] - we don't define anything - # - # Note: - # Don't mangle - user has no control over package name basename = node.name.split(".")[0] if basename == "*": # Use the ImportFrom node's line number for consistency @@ -299,9 +298,7 @@ def _get_alias_name( "is not allowed in marimo." ) return basename - else: - node.asname = self._if_local_then_mangle(node.asname) - return node.asname + return node.asname def _is_defined(self, identifier: str) -> bool: """Check if `identifier` is defined in any block.""" @@ -967,6 +964,14 @@ def visit_Name(self, node: ast.Name) -> ast.Name: ): self._add_ref(node, node.id, deleted=True) elif self.is_local(node.id): + # An unresolved underscore-prefixed Load is intentionally + # *not* added to refs. We could promote it so the dataflow + # wires it up to e.g. `from foo import _bar` in another + # cell, but that would expand the cross-cell reactive + # surface to every undefined underscore name. For now we + # accept the trade-off: cross-cell underscore imports work + # at runtime via the shared globals dict (the importer must + # have run), but the reactive graph won't track the edge. mangled_name = self._if_local_then_mangle( node.id, ignore_scope=True ) @@ -978,10 +983,12 @@ def visit_Name(self, node: ast.Name) -> ast.Name: # doesn't define it yet — this handles recursive calls to # underscore-prefixed functions, where the function name # isn't registered in the top-level block until after its - # body is visited. - if ( - block.is_defined(mangled_name) - or len(self.block_stack) > 1 + # body is visited. Skip mangling when the *unmangled* + # name is already defined at top level (e.g. by an + # underscore-prefixed import, which is never mangled). + if block.is_defined(mangled_name) or ( + len(self.block_stack) > 1 + and not block.is_defined(node.id) ): node.id = mangled_name elif block.is_defined(node.id): diff --git a/marimo/_save/hash.py b/marimo/_save/hash.py index 5638a442083..1325a3f3073 100644 --- a/marimo/_save/hash.py +++ b/marimo/_save/hash.py @@ -690,17 +690,28 @@ def serialize_and_dequeue_content_refs( imports = get_imports(scope) for local_ref in sorted(refs): ref = if_local_then_mangle(local_ref, self.cell_id) - if ref in imports: + # Imports are never mangled (see `_get_alias_name`), so the + # `_private`-style import will appear in `imports` under its + # raw name while non-import locals are still keyed by their + # mangled form. Accept either. + import_key: Name | None = ( + ref + if ref in imports + else local_ref + if local_ref in imports + else None + ) + if import_key is not None: # TODO: There may be a way to tie this in with module watching. # e.g. module watcher could mutate the version number based # last updated timestamp. version = "" module = None if self.pin_modules: - module = sys.modules[imports[ref].module] + module = sys.modules[imports[import_key].module] version = getattr(module, "__version__", "") if not version: - module = sys.modules[imports[ref].namespace] + module = sys.modules[imports[import_key].namespace] version = getattr(module, "__version__", "") content_serialization[ref] = type_sign( diff --git a/tests/_ast/test_visitor.py b/tests/_ast/test_visitor.py index 860863c7be0..a6365732e80 100644 --- a/tests/_ast/test_visitor.py +++ b/tests/_ast/test_visitor.py @@ -985,6 +985,37 @@ def test_relative_from_import() -> None: ] +def test_underscore_imports_never_mangled() -> None: + # Imports must never be mangled — the user can't control upstream + # package names. Covers `from foo import _bar`, multi-name imports, + # plain `import _foo`, and `import x as _y`. Also checks that + # underscore-imported names referenced from a nested scope + # (decorator) stay unmangled. + cases = [ + ("from foo import _bar", {"_bar"}), + ("from foo import _bar, _baz, _qux", {"_bar", "_baz", "_qux"}), + ("import _foo", {"_foo"}), + ("import marimo as _mo", {"_mo"}), + ("from a.b import _c as _d", {"_d"}), + ( + "import marimo as _private\n" + "@_private.cache\n" + "def f(x):\n" + " return _private.md('x')\n", + {"_private", "f"}, + ), + ] + for source, expected_defs in cases: + v = visitor.ScopedVisitor(mangle_prefix="cell_test") + mod = ast.parse(source) + v.visit(mod) + assert v.defs == expected_defs, source + unparsed = ast.unparse(mod) + assert "_cell_test_" not in unparsed, ( + f"import was mangled in {source!r}: {unparsed!r}" + ) + + def test_from_import_star() -> None: expr = "from a.b.c import *" v = visitor.ScopedVisitor() diff --git a/tests/_runtime/test_runtime.py b/tests/_runtime/test_runtime.py index d6cbbc711ef..8be597981ac 100644 --- a/tests/_runtime/test_runtime.py +++ b/tests/_runtime/test_runtime.py @@ -877,7 +877,9 @@ async def test_delete_nonlocal_incremental_ref_raises_name_error( async def test_import_module_as_local_var( self, any_kernel: Kernel ) -> None: - # Tests that imported names are mangled but still usable + # Imports are never mangled (see `_get_alias_name`) and are + # exposed as graph defs even when underscore-prefixed, so they + # can be referenced across cells. k = any_kernel await k.run( [ @@ -887,9 +889,50 @@ async def test_import_module_as_local_var( ), ] ) - # _sys mangled, should not be in globals - assert "_sys" not in k.globals assert k.globals["msize"] == sys.maxsize + assert "_sys" in k.graph.cells["0"].defs + + async def test_underscore_prefixed_import_in_cell( + self, any_kernel: Kernel + ) -> None: + # An underscore-prefixed `from x import _y` in a single cell + # must resolve when used in the same cell. + k = any_kernel + await k.run( + [ + ExecuteCellCommand( + cell_id="0", + code=( + "from marimo import _output\nmsg = _output.md.md('hi')" + ), + ), + ] + ) + assert not k.errors, k.errors + assert "hi" in k.globals["msg"].text + + async def test_underscore_prefixed_import_across_cells( + self, k: Kernel + ) -> None: + # Cross-cell: one cell does `from x import _y`, another uses + # `_y`. The consumer resolves at runtime via shared globals + # (the importer ran first). No reactive edge is created — the + # consumer is not in the importer's children — so editing the + # importer won't automatically re-run the consumer. We accept + # that trade-off rather than promoting every undefined + # underscore Load to a cross-cell ref. + importer = ExecuteCellCommand( + cell_id="0", + code="from marimo import _output", + ) + consumer = ExecuteCellCommand( + cell_id="1", + code="msg = _output.md.md('cross')", + ) + await k.run([importer, consumer]) + assert not k.errors, k.errors + assert "cross" in k.globals["msg"].text + assert "1" not in k.graph.children.get("0", set()) async def test_cell_transitioned_to_error_is_not_stale( self, lazy_kernel: Kernel