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

Lib/test/test_dis.py

@@ -382,36 +382,34 @@ def wrap_func_w_kwargs():
 # leading newline is for a reason (tests lineno)
 
 dis_annot_stmt_str = """\
-  --           MAKE_CELL                0 (__conditional_annotations__)
-
-   0           RESUME                   0
+  0           RESUME                   0
 
-   2           LOAD_CONST               1 (<code object __annotate__ at 0x..., file "<dis>", line 2>)
-               MAKE_FUNCTION
-               STORE_NAME               4 (__annotate__)
-               BUILD_SET                0
-               STORE_NAME               0 (__conditional_annotations__)
-               LOAD_SMALL_INT           1
-               STORE_NAME               1 (x)
-               LOAD_NAME                0 (__conditional_annotations__)
-               LOAD_SMALL_INT           0
-               CALL_INTRINSIC_2         6 (INTRINSIC_ADD_CONDITIONAL_ANNOTATION)
-               POP_TOP
+  2           LOAD_CONST               1 (<code object __annotate__ at 0x..., file "<dis>", line 2>)
+              MAKE_FUNCTION
+              STORE_NAME               4 (__annotate__)
+              BUILD_SET                0
+              STORE_NAME               0 (__conditional_annotations__)
+              LOAD_SMALL_INT           1
+              STORE_NAME               1 (x)
+              LOAD_NAME                0 (__conditional_annotations__)
+              LOAD_SMALL_INT           0
+              CALL_INTRINSIC_2         6 (INTRINSIC_ADD_CONDITIONAL_ANNOTATION)
+              POP_TOP
 
-   3           LOAD_NAME                0 (__conditional_annotations__)
-               LOAD_SMALL_INT           1
-               CALL_INTRINSIC_2         6 (INTRINSIC_ADD_CONDITIONAL_ANNOTATION)
-               POP_TOP
+  3           LOAD_NAME                0 (__conditional_annotations__)
+              LOAD_SMALL_INT           1
+              CALL_INTRINSIC_2         6 (INTRINSIC_ADD_CONDITIONAL_ANNOTATION)
+              POP_TOP
 
-   4           LOAD_SMALL_INT           1
-               LOAD_NAME                2 (lst)
-               LOAD_NAME                3 (fun)
-               PUSH_NULL
-               LOAD_SMALL_INT           0
-               CALL                     1
-               STORE_SUBSCR
-               LOAD_COMMON_CONSTANT     7 (None)
-               RETURN_VALUE
+  4           LOAD_SMALL_INT           1
+              LOAD_NAME                2 (lst)
+              LOAD_NAME                3 (fun)
+              PUSH_NULL
+              LOAD_SMALL_INT           0
+              CALL                     1
+              STORE_SUBSCR
+              LOAD_COMMON_CONSTANT     7 (None)
+              RETURN_VALUE
 """
 
 fn_with_annotate_str = """

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,60 @@ 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.assertEqual(class_code.co_cellvars,
+                         ('__classcell__', '__classdictcell__'))
+        self.assertEqual(class_code.co_freevars, ('__class__',))
+
+    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.assertEqual(class_code.co_cellvars,
+                         ('__class__', '__classcell__', '__classdictcell__'))
+
     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,36 +701,27 @@ _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;
         }
     }
-    if (u->u_ste->ste_has_conditional_annotations) {
-        /* Cook up an implicit __conditional_annotations__ cell */
-        Py_ssize_t res;
-        assert(u->u_scope_type == COMPILE_SCOPE_CLASS || u->u_scope_type == COMPILE_SCOPE_MODULE);
-        res = _PyCompile_DictAddObj(u->u_metadata.u_cellvars, &_Py_ID(__conditional_annotations__));
-        if (res < 0) {
-            compiler_unit_free(u);
-            return ERROR;
-        }
-    }
-
     u->u_metadata.u_freevars = dictbytype(u->u_ste->ste_symbols, FREE, DEF_FREE_CLASS,
                                PyDict_GET_SIZE(u->u_metadata.u_cellvars));
     if (!u->u_metadata.u_freevars) {
@@ -999,6 +990,15 @@ compiler_resolve_inlined_free(PySTEntryObject **ste, PyObject *name)
         assert(parent != NULL);
         if (parent->ste_type == ClassBlock) {
             if (_PyST_IsClassClosureName(name)) {
+                /* __conditional_annotations__ is a real CELL on the class;
+                 * other class-closure names (__class__, ...) are cooked
+                 * under private cellvar names and are not in ste_symbols. */
+                int parent_scope = _PyST_GetScope(parent, name);
+                RETURN_IF_ERROR(parent_scope);
+                if (parent_scope == CELL) {
+                    *ste = parent;
+                    return CELL;
+                }
                 return GLOBAL_IMPLICIT;
             }
             break;
@@ -1010,12 +1010,24 @@ 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__);
+    }
+    return NULL;
+}
+
 int
 _PyCompile_GetRefType(compiler *c, PyObject *name)
 {
-    if (c->u->u_scope_type == COMPILE_SCOPE_CLASS && _PyST_IsClassClosureName(name)) {
-        return CELL;
-    }
     PySTEntryObject *ste = c->u->u_ste;
     int scope = compiler_resolve_inlined_free(&ste, name);
     RETURN_IF_ERROR(scope);
@@ -1054,6 +1066,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__, ...). Map a free of the public name to that
+     * private cell, except in an inlined comprehension that already has
+     * its own CELL under the public name — prefer that slot. */
+    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

Python/symtable.c

@@ -1339,6 +1339,36 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
     else if (ste->ste_type == ClassBlock && !drop_class_free(ste, newfree)) {
         goto error;
     }
+    /* Class scopes skip analyze_cells(); ensure the deferred-annotation
+     * name is a CELL (it may appear only via drop_class_free()). Modules
+     * keep it as an implicit global in the module dict. */
+    if (ste->ste_has_conditional_annotations && ste->ste_type == ClassBlock) {
+        PyObject *name = &_Py_ID(__conditional_annotations__);
+        int contains = PyDict_Contains(ste->ste_symbols, name);
+        if (contains < 0) {
+            goto error;
+        }
+        if (!contains) {
+            PyObject *flags = PyLong_FromLong(DEF_LOCAL);
+            if (flags == NULL) {
+                goto error;
+            }
+            int rc = PyDict_SetItem(ste->ste_symbols, name, flags);
+            Py_DECREF(flags);
+            if (rc < 0) {
+                goto error;
+            }
+        }
+        PyObject *v_cell = PyLong_FromLong(CELL);
+        if (v_cell == NULL) {
+            goto error;
+        }
+        int rc = PyDict_SetItem(scopes, name, v_cell);
+        Py_DECREF(v_cell);
+        if (rc < 0) {
+            goto error;
+        }
+    }
     /* Records the results of the analysis in the symbol table entry */
     if (!update_symbols(ste->ste_symbols, scopes, bound, newfree,
                         (ste->ste_type == ClassBlock) || ste->ste_can_see_class_scope))
@@ -2879,7 +2909,12 @@ symtable_visit_annotation(struct symtable *st, expr_ty annotation, void *key)
             && !st->st_cur->ste_has_conditional_annotations)
     {
         st->st_cur->ste_has_conditional_annotations = 1;
-        if (!symtable_add_def(st, &_Py_ID(__conditional_annotations__), USE, LOCATION(annotation))) {
+        /* Class: local binding, promoted to CELL in analyze_block.
+         * Module: USE only (implicit global); the set lives in the
+         * module dict so user rebinding is visible to the intrinsic. */
+        int flag = (st->st_cur->ste_type == ClassBlock) ? DEF_LOCAL : USE;
+        if (!symtable_add_def(st, &_Py_ID(__conditional_annotations__),
+                              flag, LOCATION(annotation))) {
             return 0;
         }
     }