pythongh-114828: fix __class__ in class-scope inlined comprehensions

carljm · Feb 7, 2024 · 342f89d · 342f89d
1 parent 76108b8
commit 342f89d
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 0 deletions.
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/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_DelItem(comp->ste_symbols, k) < 0) {
+        return 0;
+    }
     return 1;
 }