gh-156664: Fix interaction between free variables and comprehensions

Lib/test/test_listcomps.py

@@ -757,6 +757,66 @@ def test_multiple_comprehension_name_reuse(self):
         self._check_in_scopes(code, {"x": 2, "y": [3]}, ns={"x": 3}, scopes=["class"])
         self._check_in_scopes(code, {"x": 2, "y": [2]}, ns={"x": 3}, scopes=["function", "module"])
 
+    def test_comprehension_name_reuse_with_free_variable(self):
+        x = 3
+
+        def sibling_comprehension():
+            [x for x in [1]]
+            return [x for _ in [1]]
+
+        self.assertEqual(sibling_comprehension(), [3])
+
+        def nested_function():
+            [x for x in [1]]
+
+            def inner():
+                return x
+
+            return inner()
+
+        self.assertEqual(nested_function(), 3)
+
+    def test_comprehension_cell_and_free_variable(self):
+        x = 3
+
+        def captured_then_sibling():
+            funcs = [lambda: x for x in [1]]
+            return funcs[0](), [x for _ in [1]]
+
+        self.assertEqual(captured_then_sibling(), (1, [3]))
+
+        def captured_then_nested_function():
+            funcs = [lambda: x for x in [1]]
+
+            def inner():
+                return x
+
+            return funcs[0](), inner()
+
+        self.assertEqual(captured_then_nested_function(), (1, 3))
+
+        def captured_then_generator_expression():
+            funcs = [lambda: x for x in [1]]
+            return funcs[0](), list(x for _ in [1])
+
+        self.assertEqual(captured_then_generator_expression(), (1, [3]))
+
+    def test_comprehension_cell_exception_cleanup(self):
+        x = 3
+
+        def raises_after_one():
+            yield 1
+            raise RuntimeError
+
+        def captured_then_exception():
+            funcs = []
+            try:
+                [funcs.append(lambda: x) for x in raises_after_one()]
+            except RuntimeError:
+                return funcs[0](), [x for _ in [1]]
+
+        self.assertEqual(captured_then_exception(), (1, [3]))
+
     def test_exception_locations(self):
         # The location of an exception raised from __init__ or
         # __next__ should be the iterator expression

Misc/NEWS.d/next/Core_and_Builtins/2026-08-30-07-36-57.gh-issue-156664.Ef7PHC.rst

@@ -0,0 +1,2 @@
+Fix scope analysis for a name bound in one inlined comprehension and used as
+a free variable by a sibling comprehension or nested function.

Python/codegen.c

@@ -4919,7 +4919,6 @@ codegen_push_inlined_comprehension_locals(compiler *c, location loc,
 {
     int in_class_block = (SYMTABLE_ENTRY(c)->ste_type == ClassBlock) &&
                           !_PyCompile_IsInInlinedComp(c);
-    PySTEntryObject *outer = SYMTABLE_ENTRY(c);
     // iterate over names bound in the comprehension and ensure we isolate
     // them from the outer scope as needed
     PyObject *k, *v;
@@ -4930,10 +4929,6 @@ codegen_push_inlined_comprehension_locals(compiler *c, location loc,
         RETURN_IF_ERROR(symbol);
         long scope = SYMBOL_TO_SCOPE(symbol);
 
-        long outsymbol = _PyST_GetSymbol(outer, k);
-        RETURN_IF_ERROR(outsymbol);
-        long outsc = SYMBOL_TO_SCOPE(outsymbol);
-
         if ((symbol & DEF_LOCAL && !(symbol & DEF_NONLOCAL)) || in_class_block) {
             // local names bound in comprehension must be isolated from
             // outer scope; push existing value (which may be NULL if
@@ -4949,11 +4944,7 @@ codegen_push_inlined_comprehension_locals(compiler *c, location loc,
             // comprehension and restore the original one after
             ADDOP_NAME(c, loc, LOAD_FAST_AND_CLEAR, k, varnames);
             if (scope == CELL) {
-                if (outsc == FREE) {
-                    ADDOP_NAME(c, loc, MAKE_CELL, k, freevars);
-                } else {
-                    ADDOP_NAME(c, loc, MAKE_CELL, k, cellvars);
-                }
+                ADDOP_NAME(c, loc, MAKE_CELL, k, cellvars);
             }
             if (PyList_Append(state->pushed_locals, k) < 0) {
                 return ERROR;

Python/compile.c

@@ -1080,12 +1080,9 @@ _PyCompile_TweakInlinedComprehensionScopes(compiler *c, location loc,
         // we need to temporarily handle it with the right scope while
         // compiling the comprehension. If it's free in the comprehension
         // scope, no special handling; it should be handled the same as the
-        // enclosing scope. (If it's free in outer scope and cell in inner
-        // scope, we can't treat it as both cell and free in the same function,
-        // but treating it as free throughout is fine; it's *_DEREF
-        // either way.)
-        if ((scope != outsc && scope != FREE && !(scope == CELL && outsc == FREE))
-                || in_class_block) {
+        // enclosing scope. A name that is a cell in the comprehension and free
+        // outside it uses separate cell and free-variable slots.
+        if ((scope != outsc && scope != FREE) || in_class_block) {
             if (state->temp_symbols == NULL) {
                 state->temp_symbols = PyDict_New();
                 if (state->temp_symbols == NULL) {

Python/symtable.c

@@ -804,7 +804,7 @@ is_free_in_any_child(PySTEntryObject *entry, PyObject *key)
 static int
 inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
                      PyObject *scopes, PyObject *comp_free,
-                     PyObject *inlined_cells)
+                     PyObject *inlined_cells, PyObject *local)
 {
     PyObject *k, *v;
     Py_ssize_t pos = 0;
@@ -880,17 +880,23 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
                 return 0;
             }
             if ((flags & DEF_BOUND) && ste->ste_type != ClassBlock) {
-                // free vars in comprehension that are locals in outer scope can
-                // now simply be locals, unless they are free in comp children,
-                // or if the outer scope is a class block
-                int ok = is_free_in_any_child(comp, k);
-                if (ok < 0) {
+                int is_local = PySet_Contains(local, k);
+                if (is_local < 0) {
                     return 0;
                 }
-                if (!ok) {
-                    if (PySet_Discard(comp_free, k) < 0) {
+                if (is_local) {
+                    // free vars in comprehension that are locals in outer scope can
+                    // now simply be locals, unless they are free in comp children,
+                    // or if the outer scope is a class block
+                    int ok = is_free_in_any_child(comp, k);
+                    if (ok < 0) {
                         return 0;
                     }
+                    if (!ok) {
+                        if (PySet_Discard(comp_free, k) < 0) {
+                            return 0;
+                        }
+                    }
                 }
             }
         }
@@ -913,32 +919,54 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
    provides the binding for the free variable.  The name should be
    marked CELL in this block and removed from the free list.
 
-   Note that the current block's free variables are included in free.
-   That's safe because no name can be free and local in the same scope.
+   Note that the current block's free variables are included in free.  A name
+   can appear local in scopes and free if the local binding was copied from an
+   inlined comprehension; such a name is not in local and must remain free.
 */
 
 static int
-analyze_cells(PyObject *scopes, PyObject *free, PyObject *inlined_cells)
+analyze_cells(PyObject *scopes, PyObject *free, PyObject *inlined_cells,
+              PyObject *local)
 {
-    PyObject *name, *v, *v_cell;
+    PyObject *name, *v, *v_cell, *v_free;
     int success = 0;
     Py_ssize_t pos = 0;
 
     v_cell = PyLong_FromLong(CELL);
     if (!v_cell)
         return 0;
+    v_free = PyLong_FromLong(FREE);
+    if (!v_free) {
+        Py_DECREF(v_cell);
+        return 0;
+    }
     while (PyDict_Next(scopes, &pos, &name, &v)) {
         long scope = PyLong_AsLong(v);
         if (scope == -1 && PyErr_Occurred()) {
             goto error;
         }
-        if (scope != LOCAL)
+        if (scope != LOCAL && scope != CELL)
             continue;
         int contains = PySet_Contains(free, name);
         if (contains < 0) {
             goto error;
         }
-        if (!contains) {
+        if (contains) {
+            int is_local = PySet_Contains(local, name);
+            if (is_local < 0) {
+                goto error;
+            }
+            if (!is_local) {
+                // This binding was copied from an inlined comprehension, not
+                // defined in this scope.  Another child may still need the
+                // name from an enclosing scope.
+                if (PyDict_SetItem(scopes, name, v_free) < 0) {
+                    goto error;
+                }
+                continue;
+            }
+        }
+        else if (scope == LOCAL) {
             contains = PySet_Contains(inlined_cells, name);
             if (contains < 0) {
                 goto error;
@@ -947,6 +975,11 @@ analyze_cells(PyObject *scopes, PyObject *free, PyObject *inlined_cells)
                 continue;
             }
         }
+        if (scope == CELL) {
+            // Retain a cell copied from an inlined comprehension if no child
+            // needs the same name as a free variable.
+            continue;
+        }
         /* Replace LOCAL with CELL for this name, and remove
            from free. It is safe to replace the value of name
            in the dict, because it will not cause a resize.
@@ -959,6 +992,7 @@ analyze_cells(PyObject *scopes, PyObject *free, PyObject *inlined_cells)
     success = 1;
  error:
     Py_DECREF(v_cell);
+    Py_DECREF(v_free);
     return success;
 }
 
@@ -1275,7 +1309,8 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
             goto error;
         }
         if (inline_comp) {
-            if (!inline_comprehension(ste, entry, scopes, child_free, inlined_cells)) {
+            if (!inline_comprehension(ste, entry, scopes, child_free,
+                                      inlined_cells, local)) {
                 Py_DECREF(child_free);
                 goto error;
             }
@@ -1303,8 +1338,11 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
     }
 
     /* Check if any local variables must be converted to cell variables */
-    if (_PyST_IsFunctionLike(ste) && !analyze_cells(scopes, newfree, inlined_cells))
+    if (_PyST_IsFunctionLike(ste) &&
+        !analyze_cells(scopes, newfree, inlined_cells, local))
+    {
         goto error;
+    }
     else if (ste->ste_type == ClassBlock && !drop_class_free(ste, newfree))
         goto error;
     /* Records the results of the analysis in the symbol table entry */