remove assumption that class scopes can't have cellvars

carljm · carljm · commit 5914d7790984 · 2023-05-05T16:28:11.000-07:00
diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py
@@ -143,6 +143,17 @@ def test_lambdas_with_free_var(self):
         outputs = {"y": [4, 4, 4, 4, 4]}
         self._check_in_scopes(code, outputs)
 
+    def test_class_scope_free_var_with_class_cell(self):
+        class C:
+            def method(self):
+                super()
+                return __class__
+            items = [(lambda: i) for i in range(5)]
+            y = [x() for x in items]
+
+        self.assertEqual(C.y, [4, 4, 4, 4, 4])
+        self.assertIs(C().method(), C)
+
     def test_inner_cell_shadows_outer(self):
         code = """
             items = [(lambda: i) for i in range(5)]
diff --git a/Python/compile.c b/Python/compile.c
@@ -1245,9 +1245,7 @@ compiler_enter_scope(struct compiler *c, identifier name,
         /* Cook up an implicit __class__ cell. */
         int res;
         assert(u->u_scope_type == COMPILER_SCOPE_CLASS);
-        assert(PyDict_GET_SIZE(u->u_metadata.u_cellvars) == 0);
-        res = PyDict_SetItem(u->u_metadata.u_cellvars, &_Py_ID(__class__),
-                             _PyLong_GetZero());
+        res = dict_add_o(u->u_metadata.u_cellvars, &_Py_ID(__class__));
         if (res < 0) {
             compiler_unit_free(u);
             return ERROR;
@@ -2245,7 +2243,6 @@ compiler_class(struct compiler *c, stmt_ty s)
                 compiler_exit_scope(c);
                 return ERROR;
             }
-            assert(i == 0);
             ADDOP_I(c, NO_LOCATION, LOAD_CLOSURE, i);
             ADDOP_I(c, NO_LOCATION, COPY, 1);
             if (compiler_nameop(c, NO_LOCATION, &_Py_ID(__classcell__), Store) < 0) {
