gh-114828: Fix __class__ in class-scope inlined comprehensions (#115139)

carljm · web-flow · commit fedbf77191ea · 2024-02-07T16:56:16.000Z
diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py
@@ -156,6 +156,18 @@ def method(self):
         self.assertEqual(C.y, [4, 4, 4, 4, 4])
         self.assertIs(C().method(), C)
 
+    def test_references_super(self):
+        code = """
+            res = [super for x in [1]]
+        """
+        self._check_in_scopes(code, outputs={"res": [super]})
+
+    def test_references___class__(self):
+        code = """
+            res = [__class__ for x in [1]]
+        """
+        self._check_in_scopes(code, raises=NameError)
+
     def test_inner_cell_shadows_outer(self):
         code = """
             items = [(lambda: i) for i in range(5)]
diff --git a/Misc/NEWS.d/next/Core and Builtins/2024-02-07-07-50-12.gh-issue-114828.nSXwMi.rst b/Misc/NEWS.d/next/Core and Builtins/2024-02-07-07-50-12.gh-issue-114828.nSXwMi.rst
@@ -0,0 +1,2 @@
+Fix compilation crashes in uncommon code examples using :func:`super` inside
+a comprehension in a class body.
diff --git a/Python/symtable.c b/Python/symtable.c
@@ -758,6 +758,8 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
 {
     PyObject *k, *v;
     Py_ssize_t pos = 0;
+    int remove_dunder_class = 0;
+
     while (PyDict_Next(comp->ste_symbols, &pos, &k, &v)) {
         // skip comprehension parameter
         long comp_flags = PyLong_AS_LONG(v);
@@ -779,6 +781,19 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
         if (!existing) {
             // name does not exist in scope, copy from comprehension
             assert(scope != FREE || PySet_Contains(comp_free, k) == 1);
+            if (scope == FREE && ste->ste_type == ClassBlock &&
+                _PyUnicode_EqualToASCIIString(k, "__class__")) {
+                // if __class__ is unbound in the enclosing class scope and free
+                // in the comprehension scope, it needs special handling; just
+                // letting it be marked as free in class scope will break due to
+                // drop_class_free
+                scope = GLOBAL_IMPLICIT;
+                only_flags &= ~DEF_FREE;
+                if (PySet_Discard(comp_free, k) < 0) {
+                    return 0;
+                }
+                remove_dunder_class = 1;
+            }
             PyObject *v_flags = PyLong_FromLong(only_flags);
             if (v_flags == NULL) {
                 return 0;
@@ -803,6 +818,10 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
             }
         }
     }
+    comp->ste_free = PySet_Size(comp_free) > 0;
+    if (remove_dunder_class && PyDict_DelItemString(comp->ste_symbols, "__class__") < 0) {
+        return 0;
+    }
     return 1;
 }
 

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	+Fix compilation crashes in uncommon code examples using :func:`super` inside
	`2`	`+a comprehension in a class body.`