[mypyc] Fix ClassVar self-references in class bodies (#21011)

Fixes mypyc/mypyc#972.

In CPython, the class body executes as a function where earlier
assignments are available to later ones e.g this is possible:

```Python3
class Foo:
  A: t.ClassVar = {1, 2, 3}
  B: t.ClassVar = {4, 5, 6}
  C: t.ClassVar = A | B
```


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.

This PR fixes this 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.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: VaggelisD

mypyc/irbuild/builder.py

@@ -233,6 +233,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_ir: ClassIR | None = None
+
         # 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

@@ -137,6 +137,13 @@ 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.
+    builder.class_body_classvars = {}
+    builder.class_body_obj = cls_builder.class_body_obj()
+    builder.class_body_ir = ir
+
     for stmt in cdef.defs.body:
         if (
             isinstance(stmt, (FuncDef, Decorator, OverloadedFuncDef))
@@ -179,13 +186,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)
 
+    # Clear class body context (nested classes are rejected above, so no need to save/restore).
+    builder.class_body_classvars = {}
+    builder.class_body_obj = None
+    builder.class_body_ir = None
+
     # Generate implicit property setters/getters
     for name, decl in ir.method_decls.items():
         if decl.implicit and decl.is_prop_getter:
@@ -232,12 +247,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)
@@ -293,6 +319,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 = 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

@@ -213,6 +213,17 @@ 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_ir is not None and builder.class_body_ir.is_ext_class:
+            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)