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

[JelleZijlstra]

https://github.com/python/cpython/commit/fcd5fb49b1d71165f3c503c3d2e74a082ddb2f21
commit: fcd5fb49b1d71165f3c503c3d2e74a082ddb2f21
branch: main
author: Carl Meyer <carl at oddbird.net>
committer: JelleZijlstra <jelle.zijlstra at gmail.com>
date: 2023-05-10T18:08:40-07:00
summary:

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

files:
M Lib/test/test_listcomps.py
M Python/symtable.c

diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py
index 92fed98dd000..1cc202bb599a 100644
--- a/Lib/test/test_listcomps.py
+++ b/Lib/test/test_listcomps.py
@@ -163,6 +163,16 @@ def test_inner_cell_shadows_outer(self):
         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 6e74d764245a..9361674bf165 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;
+                    }
                 }
             }
         }