gh-156664: Keep comprehension cells separate from free variables

Author: JelleZijlstra

Lib/test/test_listcomps.py

@@ -775,6 +775,41 @@ def inner():
 
         self.assertEqual(g(), 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 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

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

@@ -945,7 +945,7 @@ analyze_cells(PyObject *scopes, PyObject *free, PyObject *inlined_cells,
         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) {
@@ -966,7 +966,7 @@ analyze_cells(PyObject *scopes, PyObject *free, PyObject *inlined_cells,
                 continue;
             }
         }
-        else {
+        else if (scope == LOCAL) {
             contains = PySet_Contains(inlined_cells, name);
             if (contains < 0) {
                 goto error;
@@ -975,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.