Give class-closure cells private names so they don't collide with freevars

Store the class unit's synthetic cells as __classcell__/__classdictcell__
instead of __class__/__classdict__, and map method freevars of the public
names to those private cells in LookupArg. This keeps an enclosing free
named __class__ distinct from the class cell in localsplus.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: iritkatriel

Lib/test/test_super.py

@@ -4,6 +4,7 @@
 import pickle
 import textwrap
 import threading
+import types
 import unittest
 from unittest.mock import patch
 from test.support import import_helper, threading_helper
@@ -204,6 +205,63 @@ def f():
 
         self.assertIs(test_class, A)
 
+    def test_class_closure_uses_private_cellvar_name(self):
+        # The class unit's cell must not share the localsplus name __class__
+        # with a FREE load of an enclosing __class__.
+        def outer(__class__):
+            class C:
+                enclosing = __class__
+                def method(self):
+                    return __class__
+            return C
+        C = outer(object)
+        self.assertIs(C.enclosing, object)
+        self.assertIs(C().method(), C)
+
+        code = compile(
+            "def outer(__class__):\n"
+            "    class C:\n"
+            "        enclosing = __class__\n"
+            "        def method(self):\n"
+            "            return __class__\n",
+            "<test>", "exec")
+        outer_code = next(c for c in code.co_consts
+                          if isinstance(c, types.CodeType) and c.co_name == "outer")
+        class_code = next(c for c in outer_code.co_consts
+                          if isinstance(c, types.CodeType) and c.co_name == "C")
+        self.assertIn("__classcell__", class_code.co_cellvars)
+        self.assertNotIn("__class__", class_code.co_cellvars)
+        self.assertIn("__class__", class_code.co_freevars)
+        localsplus = (class_code.co_varnames + class_code.co_cellvars
+                      + class_code.co_freevars)
+        self.assertEqual(len(set(localsplus)), len(localsplus))
+
+    def test_class_cell_distinct_from_inlined_comp_cell(self):
+        # Private __classcell__ coexists with an inlined-comp cell named
+        # __class__; methods close over the private cell.
+        ns = {}
+        exec(
+            "class C:\n"
+            "    def method(self):\n"
+            "        return __class__\n"
+            "    lambdas = [lambda: __class__ for __class__ in (1, 2)]\n",
+            ns,
+        )
+        C = ns["C"]
+        self.assertIs(C().method(), C)
+        self.assertEqual([f() for f in C.lambdas], [2, 2])
+        self.assertIsNot(C.lambdas[0].__closure__[0], C.method.__closure__[0])
+
+        class_code = next(c for c in compile(
+            "class C:\n"
+            "    def method(self):\n"
+            "        return __class__\n"
+            "    lambdas = [lambda: __class__ for __class__ in (1, 2)]\n",
+            "<test>", "exec").co_consts
+            if isinstance(c, types.CodeType) and c.co_name == "C")
+        self.assertIn("__classcell__", class_code.co_cellvars)
+        self.assertIn("__class__", class_code.co_cellvars)
+
     def test___classcell___expected_behaviour(self):
         # See issue #23722
         class Meta(type):

Python/codegen.c

@@ -1641,10 +1641,9 @@ codegen_class_body(compiler *c, stmt_ty s, int firstlineno)
     if (SYMTABLE_ENTRY(c)->ste_needs_classdict) {
         ADDOP_IN_SCOPE(c, loc, LOAD_LOCALS);
 
-        // We can't use codegen_nameop here because we need to generate a
-        // STORE_DEREF in a class namespace, and codegen_nameop() won't do
-        // that by default.
-        ADDOP_N_IN_SCOPE(c, loc, STORE_DEREF, &_Py_ID(__classdict__), cellvars);
+        // Private cellvar __classdictcell__; codegen_nameop would not
+        // emit STORE_DEREF for a class-namespace nameop.
+        ADDOP_N_IN_SCOPE(c, loc, STORE_DEREF, &_Py_ID(__classdictcell__), cellvars);
     }
     if (SYMTABLE_ENTRY(c)->ste_has_conditional_annotations) {
         ADDOP_I_IN_SCOPE(c, loc, BUILD_SET, 0);
@@ -1659,7 +1658,7 @@ codegen_class_body(compiler *c, stmt_ty s, int firstlineno)
     /* Set __classdictcell__ if necessary */
     if (SYMTABLE_ENTRY(c)->ste_needs_classdict) {
         /* Store __classdictcell__ into class namespace */
-        int i = _PyCompile_LookupCellvar(c, &_Py_ID(__classdict__));
+        int i = _PyCompile_LookupCellvar(c, &_Py_ID(__classdictcell__));
         RETURN_IF_ERROR_IN_SCOPE(c, i);
         ADDOP_I_IN_SCOPE(c, NO_LOCATION, LOAD_CLOSURE, i);
         RETURN_IF_ERROR_IN_SCOPE(
@@ -1668,7 +1667,7 @@ codegen_class_body(compiler *c, stmt_ty s, int firstlineno)
     /* Return __classcell__ if it is referenced, otherwise return None */
     if (SYMTABLE_ENTRY(c)->ste_needs_class_closure) {
         /* Store __classcell__ into class namespace & return it */
-        int i = _PyCompile_LookupCellvar(c, &_Py_ID(__class__));
+        int i = _PyCompile_LookupCellvar(c, &_Py_ID(__classcell__));
         RETURN_IF_ERROR_IN_SCOPE(c, i);
         ADDOP_I_IN_SCOPE(c, NO_LOCATION, LOAD_CLOSURE, i);
         ADDOP_I_IN_SCOPE(c, NO_LOCATION, COPY, 1);

Python/compile.c

@@ -701,20 +701,22 @@ _PyCompile_EnterScope(compiler *c, identifier name, int scope_type,
         return ERROR;
     }
     if (u->u_ste->ste_needs_class_closure) {
-        /* Cook up an implicit __class__ cell. */
+        /* Private cellvar name so it cannot clash with a class-body
+         * LOCAL/FREE named __class__. Methods still expose freevar
+         * __class__; _PyCompile_LookupArg maps that free to this cell. */
         Py_ssize_t res;
         assert(u->u_scope_type == COMPILE_SCOPE_CLASS);
-        res = _PyCompile_DictAddObj(u->u_metadata.u_cellvars, &_Py_ID(__class__));
+        res = _PyCompile_DictAddObj(u->u_metadata.u_cellvars, &_Py_ID(__classcell__));
         if (res < 0) {
             compiler_unit_free(u);
             return ERROR;
         }
     }
     if (u->u_ste->ste_needs_classdict) {
-        /* Cook up an implicit __classdict__ cell. */
+        /* Private cellvar name; see ste_needs_class_closure above. */
         Py_ssize_t res;
         assert(u->u_scope_type == COMPILE_SCOPE_CLASS);
-        res = _PyCompile_DictAddObj(u->u_metadata.u_cellvars, &_Py_ID(__classdict__));
+        res = _PyCompile_DictAddObj(u->u_metadata.u_cellvars, &_Py_ID(__classdictcell__));
         if (res < 0) {
             compiler_unit_free(u);
             return ERROR;
@@ -1010,11 +1012,35 @@ compiler_resolve_inlined_free(PySTEntryObject **ste, PyObject *name)
     return scope;
 }
 
+/* Class-closure cells use private cellvar names that do not collide with
+ * class-body locals/frees of the same spelling. Method freevars keep the
+ * public names (__class__, ...); map those to the private cellvar. */
+static PyObject *
+compiler_class_closure_cellvar(PyObject *name)
+{
+    if (name == &_Py_ID(__class__)) {
+        return &_Py_ID(__classcell__);
+    }
+    if (name == &_Py_ID(__classdict__)) {
+        return &_Py_ID(__classdictcell__);
+    }
+    if (name == &_Py_ID(__conditional_annotations__)) {
+        return &_Py_ID(__conditional_annotations__);
+    }
+    return NULL;
+}
+
 int
 _PyCompile_GetRefType(compiler *c, PyObject *name)
 {
-    if (c->u->u_scope_type == COMPILE_SCOPE_CLASS && _PyST_IsClassClosureName(name)) {
-        return CELL;
+    /* Synthetic cells cooked into u_cellvars under their public spelling
+     * (today only __conditional_annotations__) are not always in ste_symbols. */
+    if (c->u->u_scope_type == COMPILE_SCOPE_CLASS ||
+        c->u->u_scope_type == COMPILE_SCOPE_MODULE) {
+        PyObject *cell_name = compiler_class_closure_cellvar(name);
+        if (cell_name != NULL && cell_name == name) {
+            return CELL;
+        }
     }
     PySTEntryObject *ste = c->u->u_ste;
     int scope = compiler_resolve_inlined_free(&ste, name);
@@ -1054,6 +1080,61 @@ _PyCompile_LookupCellvar(compiler *c, PyObject *name)
 int
 _PyCompile_LookupArg(compiler *c, PyCodeObject *co, PyObject *name)
 {
+    /* Class units store the synthetic class-closure cell under a private
+     * name (__classcell__, ...). Inside an inlined comprehension, a nested
+     * free of the public name usually captures that comprehension's CELL;
+     * elsewhere (methods, class-body lambdas) it is the private cell. */
+    if (c->u->u_scope_type == COMPILE_SCOPE_CLASS) {
+        PyObject *cell_name = compiler_class_closure_cellvar(name);
+        if (cell_name != NULL) {
+            int arg;
+            int in_inlined = (c->u->u_ste->ste_type == InlinedComprehensionBlock);
+            if (in_inlined) {
+                arg = dict_lookup_arg(c->u->u_metadata.u_cellvars, name);
+                if (arg != -1) {
+                    return arg;
+                }
+                if (PyErr_Occurred()) {
+                    return ERROR;
+                }
+            }
+            if (cell_name != name) {
+                arg = dict_lookup_arg(c->u->u_metadata.u_cellvars, cell_name);
+                if (arg != -1) {
+                    return arg;
+                }
+                if (PyErr_Occurred()) {
+                    return ERROR;
+                }
+            }
+            if (!in_inlined) {
+                arg = dict_lookup_arg(c->u->u_metadata.u_cellvars, name);
+                if (arg != -1) {
+                    return arg;
+                }
+                if (PyErr_Occurred()) {
+                    return ERROR;
+                }
+            }
+            arg = dict_lookup_arg(c->u->u_metadata.u_freevars, name);
+            if (arg != -1 || PyErr_Occurred()) {
+                return arg;
+            }
+            PyObject *freevars = _PyCode_GetFreevars(co);
+            if (freevars == NULL) {
+                PyErr_Clear();
+            }
+            PyErr_Format(PyExc_SystemError,
+                "compiler_lookup_arg(name=%R) failed in class %S; "
+                "freevars of code %S: %R",
+                name,
+                c->u->u_metadata.u_name,
+                co->co_name,
+                freevars);
+            Py_XDECREF(freevars);
+            return ERROR;
+        }
+    }
     /* Special case: If a class contains a method with a
      * free variable that has the same name as a method,
      * the name will be considered free *and* local in the