gh-104357: fix inlined comprehensions that close over iteration var (#104368)

2 files changed, 23 insertions, 6 deletions

diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py
index 92fed98dd0..1cc202bb59 100644
--- a/Lib/test/test_listcomps.py
+++ b/Lib/test/test_listcomps.py
@@ -163,6 +163,16 @@ class ListComprehensionTest(unittest.TestCase):
         outputs = {"y": [4, 4, 4, 4, 4], "i": 20}
         self._check_in_scopes(code, outputs)
 
+    def test_inner_cell_shadows_outer_no_store(self):
+        code = """
+            def f(x):
+                return [lambda: x for x in range(x)], x
+            fns, x = f(2)
+            y = [fn() for fn in fns]
+        """
+        outputs = {"y": [1, 1], "x": 2}
+        self._check_in_scopes(code, outputs)
+
     def test_closure_can_jump_over_comp_scope(self):
         code = """
             items = [(lambda: y) for i in range(5)]
diff --git a/Python/symtable.c b/Python/symtable.c
index 6e74d76424..9361674bf1 100644
--- a/Python/symtable.c
+++ b/Python/symtable.c
@@ -607,12 +607,19 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
             SET_SCOPE(scopes, k, scope);
         }
         else {
-            // free vars in comprehension that are locals in outer scope can
-            // now simply be locals, unless they are free in comp children
-            if ((PyLong_AsLong(existing) & DEF_BOUND) &&
-                    !is_free_in_any_child(comp, k)) {
-                if (PySet_Discard(comp_free, k) < 0) {
-                    return 0;
+            if (PyLong_AsLong(existing) & DEF_BOUND) {
+                // cell vars in comprehension that are locals in outer scope
+                // must be promoted to cell so u_cellvars isn't wrong
+                if (scope == CELL && ste->ste_type == FunctionBlock) {
+                    SET_SCOPE(scopes, k, scope);
+                }
+
+                // free vars in comprehension that are locals in outer scope can
+                // now simply be locals, unless they are free in comp children
+                if (!is_free_in_any_child(comp, k)) {
+                    if (PySet_Discard(comp_free, k) < 0) {
+                        return 0;
+                    }
                 }
             }
         }