fix: handle comprehension shadowing when parsing required references for variables (#7446)

This change fixes a bug in our caching mechanism's parsing of references
required for each variable. List comprehension was not catching shadowed
references.

This showed up by causing an issue in threaded caching.

Author: dmadisetti

marimo/_ast/visitor.py

@@ -1006,12 +1006,20 @@ def visit_Name(self, node: ast.Name) -> ast.Name:
 
         # Handle refs on the block scope level, or capture cell level
         # references.
+        # Only add to ref_stack if the variable is not defined in any
+        # non-module ancestor block. This prevents function parameters from
+        # being incorrectly marked as refs when used in nested scopes like
+        # list comprehensions, while still capturing module-level references.
         if (
             isinstance(node.ctx, ast.Load)
             and self._is_defined(node.id)
             and node.id not in self.ref_stack[-1]
             and (
-                node.id not in self.block_stack[-1].defs
+                # Check blocks[1:] - skip module block so module-level vars
+                # are still tracked as refs, but function params aren't
+                not any(
+                    node.id in block.defs for block in self.block_stack[1:]
+                )
                 or len(self.block_stack) == 1
             )
         ):

marimo/_save/save.py

@@ -424,7 +424,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
         try:
             if attempt.hit:
                 attempt.restore(scope)
-                return attempt.meta["return"]
+                return attempt.meta.get("return")
 
             start_time = time.time()
             response = self.__wrapped__(*args, **kwargs)

tests/_ast/test_visitor.py

@@ -305,6 +305,96 @@ def test_nested_comprehension_generator_with_named_expr() -> None:
     assert v.variable_data == {"x": [VariableData(kind="variable")]}
 
 
+def test_function_param_in_comprehension_not_required_ref() -> None:
+    """Function parameters used in list comprehensions should not be required_refs.
+
+    Regression test: The parameter `extension` was incorrectly added to required_refs
+    when used as the iterator in a list comprehension inside the function.
+    See: test_shadowed_ui_variable_threadpool in tests/_save/test_cache.py
+    """
+    code = cleandoc(
+        """
+        def helper(extension):
+            return [e for e in extension or []]
+        """
+    )
+    v = visitor.ScopedVisitor()
+    mod = ast.parse(code)
+    v.visit(mod)
+
+    assert v.defs == {"helper"}
+    assert v.refs == set()  # No external refs!
+    # extension is a PARAMETER, not an external dependency
+    # Compare to test_globals_in_functions: foo(a...) where a is not in required_refs
+    assert v.variable_data == {
+        "helper": [VariableData(kind="function", required_refs=set())]
+    }
+
+
+def test_nested_function_param_in_comprehension_not_required_ref() -> None:
+    """Ensure that additional nesting works."""
+    code = cleandoc(
+        """
+        def helper():
+            extension = []
+            def foo():
+                def bar():
+                    return [e for e in extension or []]
+                return bar
+            return foo
+        """
+    )
+    v = visitor.ScopedVisitor()
+    mod = ast.parse(code)
+    v.visit(mod)
+
+    assert v.defs == {"helper"}
+    assert v.refs == set()  # No external refs!
+    assert v.variable_data == {
+        "helper": [VariableData(kind="function", required_refs=set())]
+    }
+
+
+def test_param_in_comprehension_has_required_ref() -> None:
+    """Sanity check ref still is picked up"""
+    code = cleandoc(
+        """
+        def helper():
+            return [e for e in extension or []]
+        """
+    )
+    v = visitor.ScopedVisitor()
+    mod = ast.parse(code)
+    v.visit(mod)
+
+    assert v.defs == {"helper"}
+    assert v.refs == {"extension"}
+    assert v.variable_data == {
+        "helper": [VariableData(kind="function", required_refs={"extension"})]
+    }
+
+
+def test_shadowed_param_in_comprehension_not_required_ref() -> None:
+    """Check that a shadowed variable doesn't capture ref in module scope."""
+    code = cleandoc(
+        """
+        extension = []
+        def helper(extension):
+            return [e for e in extension or []]
+        """
+    )
+    v = visitor.ScopedVisitor()
+    mod = ast.parse(code)
+    v.visit(mod)
+
+    assert v.defs == {"helper", "extension"}
+    assert v.refs == set()
+    assert v.variable_data == {
+        "helper": [VariableData(kind="function", required_refs=set())],
+        "extension": [VariableData(kind="variable", required_refs=set())],
+    }
+
+
 def test_walrus_leaks_to_global_in_comprehension() -> None:
     code = "\n".join(
         [

tests/_save/test_cache.py

@@ -2500,6 +2500,56 @@ def g():
             assert g() == 2
             return (g, arr)
 
+    @staticmethod
+    def test_shadowed_ui_variable_threadpool(app) -> None:
+        """Test shadow error with UI-derived variable and ThreadPoolExecutor.
+
+        Bug requires:
+        1. UI element providing a cell-scoped variable (e.g. `extension`)
+        2. Helper function with same-named parameter using nested scope (list comp)
+        3. Cached function called via ThreadPoolExecutor.submit()
+
+        Causes KeyError at hash.py because scope has '*extension' (ARG_PREFIX)
+        but lookup uses 'extension'.
+        """
+        with app.setup:
+            from concurrent.futures import ThreadPoolExecutor
+
+            import marimo as mo
+
+        @app.cell
+        def _():
+            ui_input = mo.ui.text(value="hello")
+            return (ui_input,)
+
+        @app.cell
+        def _():
+            def helper(extension: list[str] | None) -> int:
+                # Nested scope using extension triggers the bug
+                # for e in extension: ... works fine.
+                return len([e for e in extension or []])
+
+            @mo.cache
+            def inner(extension: list[str] | None) -> int:
+                assert len([e for e in extension or []]) == 5
+                return helper(extension)
+
+            return (inner,)
+
+        @app.cell
+        def _(inner, ui_input):
+            extension = ui_input.value
+
+            results = []
+            # has to be in a thread submission
+            # the following works fine
+            assert inner(extension) == 5
+
+            with ThreadPoolExecutor(max_workers=2) as executor:
+                future = executor.submit(inner, extension)
+                assert future.result() == 5
+            return
+
 
 class TestPersistentCache:
     async def test_pickle_context(