Fix ClassVar self-references in class bodies

In CPython, the class body executes as a function where earlier
assignments are available to later ones (e.g. C = A | B where A
is a ClassVar defined earlier in the same class). mypyc previously
resolved such names via load_global(), looking them up in the module
globals dict where they don't exist — causing a KeyError at runtime.

Fix by tracking ClassVar names as they're defined during class body
processing, and redirecting lookups to the class being built: the
type object (py_get_attr) for extension classes, or the class dict
(dict_get_item_op) for non-extension classes.

This enables patterns like:
    class Parser:
        TYPE_TOKENS: ClassVar = {"INT", "VARCHAR"}
        FUNC_TOKENS: ClassVar = TYPE_TOKENS | {"FUNCTION"}

Author: VaggelisD

mypyc/irbuild/builder.py

@@ -228,6 +228,14 @@ def __init__(
 
         self.visitor = visitor
 
+        # Class body context: tracks ClassVar names defined so far when processing
+        # a class body, so that intra-class references (e.g. C = A | B where A is
+        # a ClassVar defined earlier in the same class) can be resolved correctly.
+        # Without this, mypyc looks up such names in module globals, which fails.
+        self.class_body_classvars: dict[str, None] = {}
+        self.class_body_obj: Value | None = None
+        self.class_body_is_ext: bool = False
+
         # This list operates similarly to a function call stack for nested functions. Whenever a
         # function definition begins to be generated, a FuncInfo instance is added to the stack,
         # and information about that function (e.g. whether it is nested, its environment class to

mypyc/irbuild/classdef.py

@@ -136,6 +136,16 @@ def transform_class_def(builder: IRBuilder, cdef: ClassDef) -> None:
     else:
         cls_builder = NonExtClassBuilder(builder, cdef)
 
+    # Set up class body context so that intra-class ClassVar references
+    # (e.g. C = A | B where A is defined earlier in the same class) can be
+    # resolved from the class being built instead of module globals.
+    saved_classvars = builder.class_body_classvars
+    saved_obj = builder.class_body_obj
+    saved_is_ext = builder.class_body_is_ext
+    builder.class_body_classvars = {}
+    builder.class_body_obj = cls_builder.class_body_obj()
+    builder.class_body_is_ext = ir.is_ext_class
+
     for stmt in cdef.defs.body:
         if (
             isinstance(stmt, (FuncDef, Decorator, OverloadedFuncDef))
@@ -178,13 +188,21 @@ def transform_class_def(builder: IRBuilder, cdef: ClassDef) -> None:
             # We want to collect class variables in a dictionary for both real
             # non-extension classes and fake dataclass ones.
             cls_builder.add_attr(lvalue, stmt)
+            # Track this ClassVar so subsequent class body statements can reference it.
+            if is_class_var(lvalue) or stmt.is_final_def:
+                builder.class_body_classvars[lvalue.name] = None
 
         elif isinstance(stmt, ExpressionStmt) and isinstance(stmt.expr, StrExpr):
             # Docstring. Ignore
             pass
         else:
             builder.error("Unsupported statement in class body", stmt.line)
 
+    # Restore previous class body context (handles nested classes).
+    builder.class_body_classvars = saved_classvars
+    builder.class_body_obj = saved_obj
+    builder.class_body_is_ext = saved_is_ext
+
     # Generate implicit property setters/getters
     for name, decl in ir.method_decls.items():
         if decl.implicit and decl.is_prop_getter:
@@ -231,12 +249,23 @@ def add_attr(self, lvalue: NameExpr, stmt: AssignmentStmt) -> None:
     def finalize(self, ir: ClassIR) -> None:
         """Perform any final operations to complete the class IR"""
 
+    def class_body_obj(self) -> Value | None:
+        """Return the object to use for loading class attributes during class body init.
+
+        For extension classes, this is the type object. For non-extension classes,
+        this is the class dict. Returns None if not applicable.
+        """
+        return None
+
 
 class NonExtClassBuilder(ClassBuilder):
     def __init__(self, builder: IRBuilder, cdef: ClassDef) -> None:
         super().__init__(builder, cdef)
         self.non_ext = self.create_non_ext_info()
 
+    def class_body_obj(self) -> Value | None:
+        return self.non_ext.dict
+
     def create_non_ext_info(self) -> NonExtClassInfo:
         non_ext_bases = populate_non_ext_bases(self.builder, self.cdef)
         non_ext_metaclass = find_non_ext_metaclass(self.builder, self.cdef, non_ext_bases)
@@ -292,6 +321,9 @@ def __init__(self, builder: IRBuilder, cdef: ClassDef) -> None:
         # If the class is not decorated, generate an extension class for it.
         self.type_obj: Value | None = allocate_class(builder, cdef)
 
+    def class_body_obj(self) -> Value | None:
+        return self.type_obj
+
     def skip_attr_default(self, name: str, stmt: AssignmentStmt) -> bool:
         """Controls whether to skip generating a default for an attribute."""
         return False

mypyc/irbuild/expression.py

@@ -198,6 +198,19 @@ def transform_name_expr(builder: IRBuilder, expr: NameExpr) -> Value:
         else:
             return builder.read(builder.get_assignment_target(expr, for_read=True), expr.line)
 
+    # If we're evaluating a class body and this name is a ClassVar defined earlier
+    # in the same class, load it from the class being built (type object for ext classes,
+    # class dict for non-ext classes) instead of module globals.
+    if builder.class_body_obj is not None and expr.name in builder.class_body_classvars:
+        if builder.class_body_is_ext:
+            return builder.py_get_attr(builder.class_body_obj, expr.name, expr.line)
+        else:
+            return builder.primitive_op(
+                dict_get_item_op,
+                [builder.class_body_obj, builder.load_str(expr.name)],
+                expr.line,
+            )
+
     return builder.load_global(expr)
 
 

mypyc/test-data/irbuild-classes.test

@@ -220,48 +220,48 @@ def __top_level__():
     r24 :: object
     r25 :: str
     r26, r27 :: object
-    r28, r29 :: bool
-    r30 :: str
-    r31 :: tuple
-    r32 :: i32
-    r33 :: bit
-    r34 :: dict
-    r35 :: str
-    r36 :: i32
-    r37 :: bit
-    r38 :: bool
-    r39 :: object
-    r40 :: str
-    r41, r42 :: object
-    r43 :: str
-    r44 :: tuple
-    r45 :: i32
-    r46 :: bit
-    r47 :: dict
-    r48 :: str
-    r49 :: i32
-    r50 :: bit
-    r51 :: bool
-    r52, r53 :: object
-    r54 :: dict
-    r55 :: str
-    r56 :: object
-    r57 :: dict
-    r58 :: str
-    r59, r60 :: object
-    r61 :: tuple
-    r62 :: str
-    r63, r64 :: object
-    r65, r66 :: bool
-    r67, r68 :: str
-    r69 :: tuple
-    r70 :: i32
-    r71 :: bit
-    r72 :: dict
-    r73 :: str
-    r74 :: i32
-    r75 :: bit
-    r76 :: bool
+    r28 :: bool
+    r29 :: str
+    r30 :: tuple
+    r31 :: i32
+    r32 :: bit
+    r33 :: dict
+    r34 :: str
+    r35 :: i32
+    r36 :: bit
+    r37 :: bool
+    r38 :: object
+    r39 :: str
+    r40, r41 :: object
+    r42 :: str
+    r43 :: tuple
+    r44 :: i32
+    r45 :: bit
+    r46 :: dict
+    r47 :: str
+    r48 :: i32
+    r49 :: bit
+    r50 :: bool
+    r51, r52 :: object
+    r53 :: dict
+    r54 :: str
+    r55 :: object
+    r56 :: dict
+    r57 :: str
+    r58, r59 :: object
+    r60 :: tuple
+    r61 :: str
+    r62, r63 :: object
+    r64 :: bool
+    r65, r66 :: str
+    r67 :: tuple
+    r68 :: i32
+    r69 :: bit
+    r70 :: dict
+    r71 :: str
+    r72 :: i32
+    r73 :: bit
+    r74 :: bool
 L0:
     r0 = builtins :: module
     r1 = load_address _Py_NoneStruct
@@ -304,51 +304,50 @@ L2:
     r31 = PyObject_SetAttr(r27, r29, r30)
     r32 = r31 >= 0 :: signed
     __main__.C = r27 :: type
-    r34 = __main__.globals :: static
-    r35 = 'C'
-    r36 = PyDict_SetItem(r34, r35, r27)
-    r37 = r36 >= 0 :: signed
-    r38 = CPy_InitSubclass(r27)
-    r39 = <error> :: object
-    r40 = '__main__'
-    r41 = __main__.S_template :: type
-    r42 = CPyType_FromTemplate(r41, r39, r40)
-    r43 = '__mypyc_attrs__'
-    r44 = CPyTuple_LoadEmptyTupleConstant()
-    r45 = PyObject_SetAttr(r42, r43, r44)
-    r46 = r45 >= 0 :: signed
-    __main__.S = r42 :: type
-    r47 = __main__.globals :: static
-    r48 = 'S'
-    r49 = PyDict_SetItem(r47, r48, r42)
-    r50 = r49 >= 0 :: signed
-    r51 = CPy_InitSubclass(r42)
-    r52 = __main__.C :: type
-    r53 = __main__.S :: type
-    r54 = __main__.globals :: static
-    r55 = 'Generic'
-    r56 = CPyDict_GetItem(r54, r55)
-    r57 = __main__.globals :: static
-    r58 = 'T'
-    r59 = CPyDict_GetItem(r57, r58)
-    r60 = PyObject_GetItem(r56, r59)
-    r61 = PyTuple_Pack(3, r52, r53, r60)
-    r62 = '__main__'
-    r63 = __main__.D_template :: type
-    r64 = CPyType_FromTemplate(r63, r61, r62)
-    r65 = D_trait_vtable_setup()
-    r66 = D_coroutine_setup(r64)
-    r67 = '__mypyc_attrs__'
-    r68 = '__dict__'
-    r69 = PyTuple_Pack(1, r68)
-    r70 = PyObject_SetAttr(r64, r67, r69)
-    r71 = r70 >= 0 :: signed
-    __main__.D = r64 :: type
-    r72 = __main__.globals :: static
-    r73 = 'D'
-    r74 = PyDict_SetItem(r72, r73, r64)
-    r75 = r74 >= 0 :: signed
-    r76 = CPy_InitSubclass(r64)
+    r33 = __main__.globals :: static
+    r34 = 'C'
+    r35 = PyDict_SetItem(r33, r34, r27)
+    r36 = r35 >= 0 :: signed
+    r37 = CPy_InitSubclass(r27)
+    r38 = <error> :: object
+    r39 = '__main__'
+    r40 = __main__.S_template :: type
+    r41 = CPyType_FromTemplate(r40, r38, r39)
+    r42 = '__mypyc_attrs__'
+    r43 = CPyTuple_LoadEmptyTupleConstant()
+    r44 = PyObject_SetAttr(r41, r42, r43)
+    r45 = r44 >= 0 :: signed
+    __main__.S = r41 :: type
+    r46 = __main__.globals :: static
+    r47 = 'S'
+    r48 = PyDict_SetItem(r46, r47, r41)
+    r49 = r48 >= 0 :: signed
+    r50 = CPy_InitSubclass(r41)
+    r51 = __main__.C :: type
+    r52 = __main__.S :: type
+    r53 = __main__.globals :: static
+    r54 = 'Generic'
+    r55 = CPyDict_GetItem(r53, r54)
+    r56 = __main__.globals :: static
+    r57 = 'T'
+    r58 = CPyDict_GetItem(r56, r57)
+    r59 = PyObject_GetItem(r55, r58)
+    r60 = PyTuple_Pack(3, r51, r52, r59)
+    r61 = '__main__'
+    r62 = __main__.D_template :: type
+    r63 = CPyType_FromTemplate(r62, r60, r61)
+    r64 = D_trait_vtable_setup()
+    r65 = '__mypyc_attrs__'
+    r66 = '__dict__'
+    r67 = PyTuple_Pack(1, r66)
+    r68 = PyObject_SetAttr(r63, r65, r67)
+    r69 = r68 >= 0 :: signed
+    __main__.D = r63 :: type
+    r70 = __main__.globals :: static
+    r71 = 'D'
+    r72 = PyDict_SetItem(r70, r71, r63)
+    r73 = r72 >= 0 :: signed
+    r74 = CPy_InitSubclass(r63)
     return 1
 
 [case testIsInstance]

mypyc/test-data/run-classes.test

@@ -5681,3 +5681,80 @@ from native import Concrete
 c = Concrete()
 assert c.value() == 42
 assert c.derived() == 42
+
+[case testClassVarSelfReference]
+# ClassVar initializers that reference other ClassVars from the same class.
+# In CPython, the class body executes as a function where earlier assignments
+# are available to later ones. mypyc must replicate this by loading from the
+# class being built (type object for ext classes, class dict for non-ext)
+# instead of module globals.
+from typing import ClassVar, Dict, Set
+
+class Ext:
+    A: ClassVar[Set[int]] = {1, 2, 3}
+    B: ClassVar[Set[int]] = {4, 5, 6}
+    C: ClassVar[Set[int]] = A | B
+
+class ExtChained:
+    X: ClassVar[Set[int]] = {1, 2}
+    Y: ClassVar[Set[int]] = X | {3}
+    Z: ClassVar[Set[int]] = Y | {4}
+
+class ExtDict:
+    BASE: ClassVar[Dict[str, int]] = {"a": 1, "b": 2}
+    EXTENDED: ClassVar[Dict[str, int]] = {**BASE, "c": 3}
+
+class ExtSub(Ext):
+    E: ClassVar[Set[int]] = {7, 8}
+
+[file driver.py]
+from native import Ext, ExtChained, ExtDict, ExtSub
+
+assert Ext.A == {1, 2, 3}
+assert Ext.B == {4, 5, 6}
+assert Ext.C == {1, 2, 3, 4, 5, 6}
+
+assert ExtChained.X == {1, 2}
+assert ExtChained.Y == {1, 2, 3}
+assert ExtChained.Z == {1, 2, 3, 4}
+
+assert ExtDict.BASE == {"a": 1, "b": 2}
+assert ExtDict.EXTENDED == {"a": 1, "b": 2, "c": 3}
+
+assert ExtSub.C == {1, 2, 3, 4, 5, 6}
+assert ExtSub.E == {7, 8}
+
+[case testClassVarSelfReferenceNonExt]
+# Same as testClassVarSelfReference but for non-extension classes
+# (e.g. decorated classes or classes with allow_interpreted_subclasses).
+from typing import ClassVar, Dict, Set
+from mypy_extensions import mypyc_attr
+
+@mypyc_attr(allow_interpreted_subclasses=True)
+class NonExt:
+    A: ClassVar[Set[str]] = {"a", "b"}
+    B: ClassVar[Set[str]] = {"c"}
+    C: ClassVar[Set[str]] = A | B
+
+@mypyc_attr(allow_interpreted_subclasses=True)
+class NonExtDict:
+    BASE: ClassVar[Dict[str, int]] = {"x": 1}
+    EXTENDED: ClassVar[Dict[str, int]] = {**BASE, "y": 2}
+
+@mypyc_attr(allow_interpreted_subclasses=True)
+class NonExtChained:
+    X: ClassVar[Set[int]] = {10}
+    Y: ClassVar[Set[int]] = X | {20}
+    Z: ClassVar[Set[int]] = Y | {30}
+
+[file driver.py]
+from native import NonExt, NonExtDict, NonExtChained
+
+assert NonExt.A == {"a", "b"}
+assert NonExt.B == {"c"}
+assert NonExt.C == {"a", "b", "c"}
+
+assert NonExtDict.BASE == {"x": 1}
+assert NonExtDict.EXTENDED == {"x": 1, "y": 2}
+
+assert NonExtChained.Z == {10, 20, 30}