Major fixes and new features

venv/lib/python3.12/site-packages/mypyc/irbuild/ast_helpers.py

@@ -0,0 +1,120 @@
+"""IRBuilder AST transform helpers shared between expressions and statements.
+
+Shared code that is tightly coupled to mypy ASTs can be put here instead of
+making mypyc.irbuild.builder larger.
+"""
+
+from __future__ import annotations
+
+from mypy.nodes import (
+    LDEF,
+    BytesExpr,
+    ComparisonExpr,
+    Expression,
+    FloatExpr,
+    IntExpr,
+    MemberExpr,
+    NameExpr,
+    OpExpr,
+    StrExpr,
+    UnaryExpr,
+    Var,
+)
+from mypyc.ir.ops import BasicBlock
+from mypyc.ir.rtypes import is_fixed_width_rtype, is_tagged
+from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.constant_fold import constant_fold_expr
+
+
+def process_conditional(
+    self: IRBuilder, e: Expression, true: BasicBlock, false: BasicBlock
+) -> None:
+    if isinstance(e, OpExpr) and e.op in ["and", "or"]:
+        if e.op == "and":
+            # Short circuit 'and' in a conditional context.
+            new = BasicBlock()
+            process_conditional(self, e.left, new, false)
+            self.activate_block(new)
+            process_conditional(self, e.right, true, false)
+        else:
+            # Short circuit 'or' in a conditional context.
+            new = BasicBlock()
+            process_conditional(self, e.left, true, new)
+            self.activate_block(new)
+            process_conditional(self, e.right, true, false)
+    elif isinstance(e, UnaryExpr) and e.op == "not":
+        process_conditional(self, e.expr, false, true)
+    else:
+        res = maybe_process_conditional_comparison(self, e, true, false)
+        if res:
+            return
+        # Catch-all for arbitrary expressions.
+        reg = self.accept(e)
+        self.add_bool_branch(reg, true, false)
+
+
+def maybe_process_conditional_comparison(
+    self: IRBuilder, e: Expression, true: BasicBlock, false: BasicBlock
+) -> bool:
+    """Transform simple tagged integer comparisons in a conditional context.
+
+    Return True if the operation is supported (and was transformed). Otherwise,
+    do nothing and return False.
+
+    Args:
+        e: Arbitrary expression
+        true: Branch target if comparison is true
+        false: Branch target if comparison is false
+    """
+    if not isinstance(e, ComparisonExpr) or len(e.operands) != 2:
+        return False
+    ltype = self.node_type(e.operands[0])
+    rtype = self.node_type(e.operands[1])
+    if not (
+        (is_tagged(ltype) or is_fixed_width_rtype(ltype))
+        and (is_tagged(rtype) or is_fixed_width_rtype(rtype))
+    ):
+        return False
+    op = e.operators[0]
+    if op not in ("==", "!=", "<", "<=", ">", ">="):
+        return False
+    left_expr = e.operands[0]
+    right_expr = e.operands[1]
+    borrow_left = is_borrow_friendly_expr(self, right_expr)
+    left = self.accept(left_expr, can_borrow=borrow_left)
+    right = self.accept(right_expr, can_borrow=True)
+    if is_fixed_width_rtype(ltype) or is_fixed_width_rtype(rtype):
+        if not is_fixed_width_rtype(ltype):
+            left = self.coerce(left, rtype, e.line)
+        elif not is_fixed_width_rtype(rtype):
+            right = self.coerce(right, ltype, e.line)
+        reg = self.binary_op(left, right, op, e.line)
+        self.builder.flush_keep_alives()
+        self.add_bool_branch(reg, true, false)
+    else:
+        # "left op right" for two tagged integers
+        self.builder.compare_tagged_condition(left, right, op, true, false, e.line)
+    return True
+
+
+def is_borrow_friendly_expr(self: IRBuilder, expr: Expression) -> bool:
+    """Can the result of the expression borrowed temporarily?
+
+    Borrowing means keeping a reference without incrementing the reference count.
+    """
+    if isinstance(expr, (IntExpr, FloatExpr, StrExpr, BytesExpr)):
+        # Literals are immortal and can always be borrowed
+        return True
+    if (
+        isinstance(expr, (UnaryExpr, OpExpr, NameExpr, MemberExpr))
+        and constant_fold_expr(self, expr) is not None
+    ):
+        # Literal expressions are similar to literals
+        return True
+    if isinstance(expr, NameExpr):
+        if isinstance(expr.node, Var) and expr.kind == LDEF:
+            # Local variable reference can be borrowed
+            return True
+    if isinstance(expr, MemberExpr) and self.is_native_attr_ref(expr):
+        return True
+    return False

venv/lib/python3.12/site-packages/mypyc/irbuild/callable_class.py

@@ -0,0 +1,173 @@
+"""Generate a class that represents a nested function.
+
+The class defines __call__ for calling the function and allows access to
+non-local variables defined in outer scopes.
+"""
+
+from __future__ import annotations
+
+from mypyc.common import ENV_ATTR_NAME, SELF_NAME
+from mypyc.ir.class_ir import ClassIR
+from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature, RuntimeArg
+from mypyc.ir.ops import BasicBlock, Call, Register, Return, SetAttr, Value
+from mypyc.ir.rtypes import RInstance, object_rprimitive
+from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.context import FuncInfo, ImplicitClass
+from mypyc.primitives.misc_ops import method_new_op
+
+
+def setup_callable_class(builder: IRBuilder) -> None:
+    """Generate an (incomplete) callable class representing a function.
+
+    This can be a nested function or a function within a non-extension
+    class.  Also set up the 'self' variable for that class.
+
+    This takes the most recently visited function and returns a
+    ClassIR to represent that function. Each callable class contains
+    an environment attribute which points to another ClassIR
+    representing the environment class where some of its variables can
+    be accessed.
+
+    Note that some methods, such as '__call__', are not yet
+    created here. Use additional functions, such as
+    add_call_to_callable_class(), to add them.
+
+    Return a newly constructed ClassIR representing the callable
+    class for the nested function.
+    """
+    # Check to see that the name has not already been taken. If so,
+    # rename the class. We allow multiple uses of the same function
+    # name because this is valid in if-else blocks. Example:
+    #
+    #     if True:
+    #         def foo():          ---->    foo_obj()
+    #             return True
+    #     else:
+    #         def foo():          ---->    foo_obj_0()
+    #             return False
+    name = base_name = f"{builder.fn_info.namespaced_name()}_obj"
+    count = 0
+    while name in builder.callable_class_names:
+        name = base_name + "_" + str(count)
+        count += 1
+    builder.callable_class_names.add(name)
+
+    # Define the actual callable class ClassIR, and set its
+    # environment to point at the previously defined environment
+    # class.
+    callable_class_ir = ClassIR(name, builder.module_name, is_generated=True)
+
+    # The functools @wraps decorator attempts to call setattr on
+    # nested functions, so we create a dict for these nested
+    # functions.
+    # https://github.com/python/cpython/blob/3.7/Lib/functools.py#L58
+    if builder.fn_info.is_nested:
+        callable_class_ir.has_dict = True
+
+    # If the enclosing class doesn't contain nested (which will happen if
+    # this is a toplevel lambda), don't set up an environment.
+    if builder.fn_infos[-2].contains_nested:
+        callable_class_ir.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_infos[-2].env_class)
+    callable_class_ir.mro = [callable_class_ir]
+    builder.fn_info.callable_class = ImplicitClass(callable_class_ir)
+    builder.classes.append(callable_class_ir)
+
+    # Add a 'self' variable to the environment of the callable class,
+    # and store that variable in a register to be accessed later.
+    self_target = builder.add_self_to_env(callable_class_ir)
+    builder.fn_info.callable_class.self_reg = builder.read(self_target, builder.fn_info.fitem.line)
+
+
+def add_call_to_callable_class(
+    builder: IRBuilder,
+    args: list[Register],
+    blocks: list[BasicBlock],
+    sig: FuncSignature,
+    fn_info: FuncInfo,
+) -> FuncIR:
+    """Generate a '__call__' method for a callable class representing a nested function.
+
+    This takes the blocks and signature associated with a function
+    definition and uses those to build the '__call__' method of a
+    given callable class, used to represent that function.
+    """
+    # Since we create a method, we also add a 'self' parameter.
+    nargs = len(sig.args) - sig.num_bitmap_args
+    sig = FuncSignature(
+        (RuntimeArg(SELF_NAME, object_rprimitive),) + sig.args[:nargs], sig.ret_type
+    )
+    call_fn_decl = FuncDecl("__call__", fn_info.callable_class.ir.name, builder.module_name, sig)
+    call_fn_ir = FuncIR(
+        call_fn_decl, args, blocks, fn_info.fitem.line, traceback_name=fn_info.fitem.name
+    )
+    fn_info.callable_class.ir.methods["__call__"] = call_fn_ir
+    fn_info.callable_class.ir.method_decls["__call__"] = call_fn_decl
+    return call_fn_ir
+
+
+def add_get_to_callable_class(builder: IRBuilder, fn_info: FuncInfo) -> None:
+    """Generate the '__get__' method for a callable class."""
+    line = fn_info.fitem.line
+    with builder.enter_method(
+        fn_info.callable_class.ir,
+        "__get__",
+        object_rprimitive,
+        fn_info,
+        self_type=object_rprimitive,
+    ):
+        instance = builder.add_argument("instance", object_rprimitive)
+        builder.add_argument("owner", object_rprimitive)
+
+        # If accessed through the class, just return the callable
+        # object. If accessed through an object, create a new bound
+        # instance method object.
+        instance_block, class_block = BasicBlock(), BasicBlock()
+        comparison = builder.translate_is_op(
+            builder.read(instance), builder.none_object(), "is", line
+        )
+        builder.add_bool_branch(comparison, class_block, instance_block)
+
+        builder.activate_block(class_block)
+        builder.add(Return(builder.self()))
+
+        builder.activate_block(instance_block)
+        builder.add(
+            Return(builder.call_c(method_new_op, [builder.self(), builder.read(instance)], line))
+        )
+
+
+def instantiate_callable_class(builder: IRBuilder, fn_info: FuncInfo) -> Value:
+    """Create an instance of a callable class for a function.
+
+    Calls to the function will actually call this instance.
+
+    Note that fn_info refers to the function being assigned, whereas
+    builder.fn_info refers to the function encapsulating the function
+    being turned into a callable class.
+    """
+    fitem = fn_info.fitem
+    func_reg = builder.add(Call(fn_info.callable_class.ir.ctor, [], fitem.line))
+
+    # Set the environment attribute of the callable class to point at
+    # the environment class defined in the callable class' immediate
+    # outer scope. Note that there are three possible environment
+    # class registers we may use. This depends on what the encapsulating
+    # (parent) function is:
+    #
+    # - A nested function: the callable class is instantiated
+    #   from the current callable class' '__call__' function, and hence
+    #   the callable class' environment register is used.
+    # - A generator function: the callable class is instantiated
+    #   from the '__next__' method of the generator class, and hence the
+    #   environment of the generator class is used.
+    # - Regular function: we use the environment of the original function.
+    curr_env_reg = None
+    if builder.fn_info.is_generator:
+        curr_env_reg = builder.fn_info.generator_class.curr_env_reg
+    elif builder.fn_info.is_nested:
+        curr_env_reg = builder.fn_info.callable_class.curr_env_reg
+    elif builder.fn_info.contains_nested:
+        curr_env_reg = builder.fn_info.curr_env_reg
+    if curr_env_reg:
+        builder.add(SetAttr(func_reg, ENV_ATTR_NAME, curr_env_reg, fitem.line))
+    return func_reg

venv/lib/python3.12/site-packages/mypyc/irbuild/classdef.py

@@ -0,0 +1,850 @@
+"""Transform class definitions from the mypy AST form to IR."""
+
+from __future__ import annotations
+
+import typing_extensions
+from abc import abstractmethod
+from typing import Callable, Final
+
+from mypy.nodes import (
+    AssignmentStmt,
+    CallExpr,
+    ClassDef,
+    Decorator,
+    ExpressionStmt,
+    FuncDef,
+    Lvalue,
+    MemberExpr,
+    NameExpr,
+    OverloadedFuncDef,
+    PassStmt,
+    RefExpr,
+    StrExpr,
+    TempNode,
+    TypeInfo,
+    is_class_var,
+)
+from mypy.types import ENUM_REMOVED_PROPS, Instance, UnboundType, get_proper_type
+from mypyc.common import PROPSET_PREFIX
+from mypyc.ir.class_ir import ClassIR, NonExtClassInfo
+from mypyc.ir.func_ir import FuncDecl, FuncSignature
+from mypyc.ir.ops import (
+    NAMESPACE_TYPE,
+    BasicBlock,
+    Branch,
+    Call,
+    InitStatic,
+    LoadAddress,
+    LoadErrorValue,
+    LoadStatic,
+    MethodCall,
+    Register,
+    Return,
+    SetAttr,
+    TupleSet,
+    Value,
+)
+from mypyc.ir.rtypes import (
+    RType,
+    bool_rprimitive,
+    dict_rprimitive,
+    is_none_rprimitive,
+    is_object_rprimitive,
+    is_optional_type,
+    object_rprimitive,
+)
+from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.function import (
+    gen_property_getter_ir,
+    gen_property_setter_ir,
+    handle_ext_method,
+    handle_non_ext_method,
+    load_type,
+)
+from mypyc.irbuild.util import dataclass_type, get_func_def, is_constant, is_dataclass_decorator
+from mypyc.primitives.dict_ops import dict_new_op, dict_set_item_op
+from mypyc.primitives.generic_ops import py_hasattr_op, py_setattr_op
+from mypyc.primitives.misc_ops import (
+    dataclass_sleight_of_hand,
+    not_implemented_op,
+    py_calc_meta_op,
+    pytype_from_template_op,
+    type_object_op,
+)
+
+
+def transform_class_def(builder: IRBuilder, cdef: ClassDef) -> None:
+    """Create IR for a class definition.
+
+    This can generate both extension (native) and non-extension
+    classes.  These are generated in very different ways. In the
+    latter case we construct a Python type object at runtime by doing
+    the equivalent of "type(name, bases, dict)" in IR. Extension
+    classes are defined via C structs that are generated later in
+    mypyc.codegen.emitclass.
+
+    This is the main entry point to this module.
+    """
+    ir = builder.mapper.type_to_ir[cdef.info]
+
+    # We do this check here because the base field of parent
+    # classes aren't necessarily populated yet at
+    # prepare_class_def time.
+    if any(ir.base_mro[i].base != ir.base_mro[i + 1] for i in range(len(ir.base_mro) - 1)):
+        builder.error("Multiple inheritance is not supported (except for traits)", cdef.line)
+
+    if ir.allow_interpreted_subclasses:
+        for parent in ir.mro:
+            if not parent.allow_interpreted_subclasses:
+                builder.error(
+                    'Base class "{}" does not allow interpreted subclasses'.format(
+                        parent.fullname
+                    ),
+                    cdef.line,
+                )
+
+    # Currently, we only create non-extension classes for classes that are
+    # decorated or inherit from Enum. Classes decorated with @trait do not
+    # apply here, and are handled in a different way.
+    if ir.is_ext_class:
+        cls_type = dataclass_type(cdef)
+        if cls_type is None:
+            cls_builder: ClassBuilder = ExtClassBuilder(builder, cdef)
+        elif cls_type in ["dataclasses", "attr-auto"]:
+            cls_builder = DataClassBuilder(builder, cdef)
+        elif cls_type == "attr":
+            cls_builder = AttrsClassBuilder(builder, cdef)
+        else:
+            raise ValueError(cls_type)
+    else:
+        cls_builder = NonExtClassBuilder(builder, cdef)
+
+    for stmt in cdef.defs.body:
+        if isinstance(stmt, OverloadedFuncDef) and stmt.is_property:
+            if isinstance(cls_builder, NonExtClassBuilder):
+                # properties with both getters and setters in non_extension
+                # classes not supported
+                builder.error("Property setters not supported in non-extension classes", stmt.line)
+            for item in stmt.items:
+                with builder.catch_errors(stmt.line):
+                    cls_builder.add_method(get_func_def(item))
+        elif isinstance(stmt, (FuncDef, Decorator, OverloadedFuncDef)):
+            # Ignore plugin generated methods (since they have no
+            # bodies to compile and will need to have the bodies
+            # provided by some other mechanism.)
+            if cdef.info.names[stmt.name].plugin_generated:
+                continue
+            with builder.catch_errors(stmt.line):
+                cls_builder.add_method(get_func_def(stmt))
+        elif isinstance(stmt, PassStmt):
+            continue
+        elif isinstance(stmt, AssignmentStmt):
+            if len(stmt.lvalues) != 1:
+                builder.error("Multiple assignment in class bodies not supported", stmt.line)
+                continue
+            lvalue = stmt.lvalues[0]
+            if not isinstance(lvalue, NameExpr):
+                builder.error(
+                    "Only assignment to variables is supported in class bodies", stmt.line
+                )
+                continue
+            # 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)
+
+        elif isinstance(stmt, ExpressionStmt) and isinstance(stmt.expr, StrExpr):
+            # Docstring. Ignore
+            pass
+        else:
+            builder.error("Unsupported statement in class body", stmt.line)
+
+    # Generate implicit property setters/getters
+    for name, decl in ir.method_decls.items():
+        if decl.implicit and decl.is_prop_getter:
+            getter_ir = gen_property_getter_ir(builder, decl, cdef, ir.is_trait)
+            builder.functions.append(getter_ir)
+            ir.methods[getter_ir.decl.name] = getter_ir
+
+            setter_ir = None
+            setter_name = PROPSET_PREFIX + name
+            if setter_name in ir.method_decls:
+                setter_ir = gen_property_setter_ir(
+                    builder, ir.method_decls[setter_name], cdef, ir.is_trait
+                )
+                builder.functions.append(setter_ir)
+                ir.methods[setter_name] = setter_ir
+
+            ir.properties[name] = (getter_ir, setter_ir)
+            # TODO: Generate glue method if needed?
+            # TODO: Do we need interpreted glue methods? Maybe not?
+
+    cls_builder.finalize(ir)
+
+
+class ClassBuilder:
+    """Create IR for a class definition.
+
+    This is an abstract base class.
+    """
+
+    def __init__(self, builder: IRBuilder, cdef: ClassDef) -> None:
+        self.builder = builder
+        self.cdef = cdef
+        self.attrs_to_cache: list[tuple[Lvalue, RType]] = []
+
+    @abstractmethod
+    def add_method(self, fdef: FuncDef) -> None:
+        """Add a method to the class IR"""
+
+    @abstractmethod
+    def add_attr(self, lvalue: NameExpr, stmt: AssignmentStmt) -> None:
+        """Add an attribute to the class IR"""
+
+    @abstractmethod
+    def finalize(self, ir: ClassIR) -> None:
+        """Perform any final operations to complete the class IR"""
+
+
+class NonExtClassBuilder(ClassBuilder):
+    def __init__(self, builder: IRBuilder, cdef: ClassDef) -> None:
+        super().__init__(builder, cdef)
+        self.non_ext = self.create_non_ext_info()
+
+    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)
+        non_ext_dict = setup_non_ext_dict(
+            self.builder, self.cdef, non_ext_metaclass, non_ext_bases
+        )
+        # We populate __annotations__ for non-extension classes
+        # because dataclasses uses it to determine which attributes to compute on.
+        # TODO: Maybe generate more precise types for annotations
+        non_ext_anns = self.builder.call_c(dict_new_op, [], self.cdef.line)
+        return NonExtClassInfo(non_ext_dict, non_ext_bases, non_ext_anns, non_ext_metaclass)
+
+    def add_method(self, fdef: FuncDef) -> None:
+        handle_non_ext_method(self.builder, self.non_ext, self.cdef, fdef)
+
+    def add_attr(self, lvalue: NameExpr, stmt: AssignmentStmt) -> None:
+        add_non_ext_class_attr_ann(self.builder, self.non_ext, lvalue, stmt)
+        add_non_ext_class_attr(
+            self.builder, self.non_ext, lvalue, stmt, self.cdef, self.attrs_to_cache
+        )
+
+    def finalize(self, ir: ClassIR) -> None:
+        # Dynamically create the class via the type constructor
+        non_ext_class = load_non_ext_class(self.builder, ir, self.non_ext, self.cdef.line)
+        non_ext_class = load_decorated_class(self.builder, self.cdef, non_ext_class)
+
+        # Save the decorated class
+        self.builder.add(
+            InitStatic(non_ext_class, self.cdef.name, self.builder.module_name, NAMESPACE_TYPE)
+        )
+
+        # Add the non-extension class to the dict
+        self.builder.call_c(
+            dict_set_item_op,
+            [
+                self.builder.load_globals_dict(),
+                self.builder.load_str(self.cdef.name),
+                non_ext_class,
+            ],
+            self.cdef.line,
+        )
+
+        # Cache any cacheable class attributes
+        cache_class_attrs(self.builder, self.attrs_to_cache, self.cdef)
+
+
+class ExtClassBuilder(ClassBuilder):
+    def __init__(self, builder: IRBuilder, cdef: ClassDef) -> None:
+        super().__init__(builder, cdef)
+        # If the class is not decorated, generate an extension class for it.
+        self.type_obj: Value | None = allocate_class(builder, cdef)
+
+    def skip_attr_default(self, name: str, stmt: AssignmentStmt) -> bool:
+        """Controls whether to skip generating a default for an attribute."""
+        return False
+
+    def add_method(self, fdef: FuncDef) -> None:
+        handle_ext_method(self.builder, self.cdef, fdef)
+
+    def add_attr(self, lvalue: NameExpr, stmt: AssignmentStmt) -> None:
+        # Variable declaration with no body
+        if isinstance(stmt.rvalue, TempNode):
+            return
+        # Only treat marked class variables as class variables.
+        if not (is_class_var(lvalue) or stmt.is_final_def):
+            return
+        typ = self.builder.load_native_type_object(self.cdef.fullname)
+        value = self.builder.accept(stmt.rvalue)
+        self.builder.call_c(
+            py_setattr_op, [typ, self.builder.load_str(lvalue.name), value], stmt.line
+        )
+        if self.builder.non_function_scope() and stmt.is_final_def:
+            self.builder.init_final_static(lvalue, value, self.cdef.name)
+
+    def finalize(self, ir: ClassIR) -> None:
+        attrs_with_defaults, default_assignments = find_attr_initializers(
+            self.builder, self.cdef, self.skip_attr_default
+        )
+        ir.attrs_with_defaults.update(attrs_with_defaults)
+        generate_attr_defaults_init(self.builder, self.cdef, default_assignments)
+        create_ne_from_eq(self.builder, self.cdef)
+
+
+class DataClassBuilder(ExtClassBuilder):
+    # controls whether an __annotations__ attribute should be added to the class
+    # __dict__.  This is not desirable for attrs classes where auto_attribs is
+    # disabled, as attrs will reject it.
+    add_annotations_to_dict = True
+
+    def __init__(self, builder: IRBuilder, cdef: ClassDef) -> None:
+        super().__init__(builder, cdef)
+        self.non_ext = self.create_non_ext_info()
+
+    def create_non_ext_info(self) -> NonExtClassInfo:
+        """Set up a NonExtClassInfo to track dataclass attributes.
+
+        In addition to setting up a normal extension class for dataclasses,
+        we also collect its class attributes like a non-extension class so
+        that we can hand them to the dataclass decorator.
+        """
+        return NonExtClassInfo(
+            self.builder.call_c(dict_new_op, [], self.cdef.line),
+            self.builder.add(TupleSet([], self.cdef.line)),
+            self.builder.call_c(dict_new_op, [], self.cdef.line),
+            self.builder.add(LoadAddress(type_object_op.type, type_object_op.src, self.cdef.line)),
+        )
+
+    def skip_attr_default(self, name: str, stmt: AssignmentStmt) -> bool:
+        return stmt.type is not None
+
+    def get_type_annotation(self, stmt: AssignmentStmt) -> TypeInfo | None:
+        # We populate __annotations__ because dataclasses uses it to determine
+        # which attributes to compute on.
+        ann_type = get_proper_type(stmt.type)
+        if isinstance(ann_type, Instance):
+            return ann_type.type
+        return None
+
+    def add_attr(self, lvalue: NameExpr, stmt: AssignmentStmt) -> None:
+        add_non_ext_class_attr_ann(
+            self.builder, self.non_ext, lvalue, stmt, self.get_type_annotation
+        )
+        add_non_ext_class_attr(
+            self.builder, self.non_ext, lvalue, stmt, self.cdef, self.attrs_to_cache
+        )
+        super().add_attr(lvalue, stmt)
+
+    def finalize(self, ir: ClassIR) -> None:
+        """Generate code to finish instantiating a dataclass.
+
+        This works by replacing all of the attributes on the class
+        (which will be descriptors) with whatever they would be in a
+        non-extension class, calling dataclass, then switching them back.
+
+        The resulting class is an extension class and instances of it do not
+        have a __dict__ (unless something else requires it).
+        All methods written explicitly in the source are compiled and
+        may be called through the vtable while the methods generated
+        by dataclasses are interpreted and may not be.
+
+        (If we just called dataclass without doing this, it would think that all
+        of the descriptors for our attributes are default values and generate an
+        incorrect constructor. We need to do the switch so that dataclass gets the
+        appropriate defaults.)
+        """
+        super().finalize(ir)
+        assert self.type_obj
+        add_dunders_to_non_ext_dict(
+            self.builder, self.non_ext, self.cdef.line, self.add_annotations_to_dict
+        )
+        dec = self.builder.accept(
+            next(d for d in self.cdef.decorators if is_dataclass_decorator(d))
+        )
+        self.builder.call_c(
+            dataclass_sleight_of_hand,
+            [dec, self.type_obj, self.non_ext.dict, self.non_ext.anns],
+            self.cdef.line,
+        )
+
+
+class AttrsClassBuilder(DataClassBuilder):
+    """Create IR for an attrs class where auto_attribs=False (the default).
+
+    When auto_attribs is enabled, attrs classes behave similarly to dataclasses
+    (i.e. types are stored as annotations on the class) and are thus handled
+    by DataClassBuilder, but when auto_attribs is disabled the types are
+    provided via attr.ib(type=...)
+    """
+
+    add_annotations_to_dict = False
+
+    def skip_attr_default(self, name: str, stmt: AssignmentStmt) -> bool:
+        return True
+
+    def get_type_annotation(self, stmt: AssignmentStmt) -> TypeInfo | None:
+        if isinstance(stmt.rvalue, CallExpr):
+            # find the type arg in `attr.ib(type=str)`
+            callee = stmt.rvalue.callee
+            if (
+                isinstance(callee, MemberExpr)
+                and callee.fullname in ["attr.ib", "attr.attr"]
+                and "type" in stmt.rvalue.arg_names
+            ):
+                index = stmt.rvalue.arg_names.index("type")
+                type_name = stmt.rvalue.args[index]
+                if isinstance(type_name, NameExpr) and isinstance(type_name.node, TypeInfo):
+                    lvalue = stmt.lvalues[0]
+                    assert isinstance(lvalue, NameExpr)
+                    return type_name.node
+        return None
+
+
+def allocate_class(builder: IRBuilder, cdef: ClassDef) -> Value:
+    # OK AND NOW THE FUN PART
+    base_exprs = cdef.base_type_exprs + cdef.removed_base_type_exprs
+    if base_exprs:
+        bases = [builder.accept(x) for x in base_exprs]
+        tp_bases = builder.new_tuple(bases, cdef.line)
+    else:
+        tp_bases = builder.add(LoadErrorValue(object_rprimitive, is_borrowed=True))
+    modname = builder.load_str(builder.module_name)
+    template = builder.add(
+        LoadStatic(object_rprimitive, cdef.name + "_template", builder.module_name, NAMESPACE_TYPE)
+    )
+    # Create the class
+    tp = builder.call_c(pytype_from_template_op, [template, tp_bases, modname], cdef.line)
+    # Immediately fix up the trait vtables, before doing anything with the class.
+    ir = builder.mapper.type_to_ir[cdef.info]
+    if not ir.is_trait and not ir.builtin_base:
+        builder.add(
+            Call(
+                FuncDecl(
+                    cdef.name + "_trait_vtable_setup",
+                    None,
+                    builder.module_name,
+                    FuncSignature([], bool_rprimitive),
+                ),
+                [],
+                -1,
+            )
+        )
+    # Populate a '__mypyc_attrs__' field containing the list of attrs
+    builder.call_c(
+        py_setattr_op,
+        [
+            tp,
+            builder.load_str("__mypyc_attrs__"),
+            create_mypyc_attrs_tuple(builder, builder.mapper.type_to_ir[cdef.info], cdef.line),
+        ],
+        cdef.line,
+    )
+
+    # Save the class
+    builder.add(InitStatic(tp, cdef.name, builder.module_name, NAMESPACE_TYPE))
+
+    # Add it to the dict
+    builder.call_c(
+        dict_set_item_op, [builder.load_globals_dict(), builder.load_str(cdef.name), tp], cdef.line
+    )
+
+    return tp
+
+
+# Mypy uses these internally as base classes of TypedDict classes. These are
+# lies and don't have any runtime equivalent.
+MAGIC_TYPED_DICT_CLASSES: Final[tuple[str, ...]] = (
+    "typing._TypedDict",
+    "typing_extensions._TypedDict",
+)
+
+
+def populate_non_ext_bases(builder: IRBuilder, cdef: ClassDef) -> Value:
+    """Create base class tuple of a non-extension class.
+
+    The tuple is passed to the metaclass constructor.
+    """
+    is_named_tuple = cdef.info.is_named_tuple
+    ir = builder.mapper.type_to_ir[cdef.info]
+    bases = []
+    for cls in cdef.info.mro[1:]:
+        if cls.fullname == "builtins.object":
+            continue
+        if is_named_tuple and cls.fullname in (
+            "typing.Sequence",
+            "typing.Iterable",
+            "typing.Collection",
+            "typing.Reversible",
+            "typing.Container",
+            "typing.Sized",
+        ):
+            # HAX: Synthesized base classes added by mypy don't exist at runtime, so skip them.
+            #      This could break if they were added explicitly, though...
+            continue
+        # Add the current class to the base classes list of concrete subclasses
+        if cls in builder.mapper.type_to_ir:
+            base_ir = builder.mapper.type_to_ir[cls]
+            if base_ir.children is not None:
+                base_ir.children.append(ir)
+
+        if cls.fullname in MAGIC_TYPED_DICT_CLASSES:
+            # HAX: Mypy internally represents TypedDict classes differently from what
+            #      should happen at runtime. Replace with something that works.
+            module = "typing"
+            if builder.options.capi_version < (3, 9):
+                name = "TypedDict"
+                if builder.options.capi_version < (3, 8):
+                    # TypedDict was added to typing in Python 3.8.
+                    module = "typing_extensions"
+                    # TypedDict is not a real type on typing_extensions 4.7.0+
+                    name = "_TypedDict"
+                    if isinstance(typing_extensions.TypedDict, type):
+                        raise RuntimeError(
+                            "It looks like you may have an old version "
+                            "of typing_extensions installed. "
+                            "typing_extensions>=4.7.0 is required on Python 3.7."
+                        )
+            else:
+                # In Python 3.9 TypedDict is not a real type.
+                name = "_TypedDict"
+            base = builder.get_module_attr(module, name, cdef.line)
+        elif is_named_tuple and cls.fullname == "builtins.tuple":
+            if builder.options.capi_version < (3, 9):
+                name = "NamedTuple"
+            else:
+                # This was changed in Python 3.9.
+                name = "_NamedTuple"
+            base = builder.get_module_attr("typing", name, cdef.line)
+        else:
+            cls_module = cls.fullname.rsplit(".", 1)[0]
+            if cls_module == builder.current_module:
+                base = builder.load_global_str(cls.name, cdef.line)
+            else:
+                base = builder.load_module_attr_by_fullname(cls.fullname, cdef.line)
+        bases.append(base)
+        if cls.fullname in MAGIC_TYPED_DICT_CLASSES:
+            # The remaining base classes are synthesized by mypy and should be ignored.
+            break
+    return builder.new_tuple(bases, cdef.line)
+
+
+def find_non_ext_metaclass(builder: IRBuilder, cdef: ClassDef, bases: Value) -> Value:
+    """Find the metaclass of a class from its defs and bases."""
+    if cdef.metaclass:
+        declared_metaclass = builder.accept(cdef.metaclass)
+    else:
+        if cdef.info.typeddict_type is not None and builder.options.capi_version >= (3, 9):
+            # In Python 3.9, the metaclass for class-based TypedDict is typing._TypedDictMeta.
+            # We can't easily calculate it generically, so special case it.
+            return builder.get_module_attr("typing", "_TypedDictMeta", cdef.line)
+        elif cdef.info.is_named_tuple and builder.options.capi_version >= (3, 9):
+            # In Python 3.9, the metaclass for class-based NamedTuple is typing.NamedTupleMeta.
+            # We can't easily calculate it generically, so special case it.
+            return builder.get_module_attr("typing", "NamedTupleMeta", cdef.line)
+
+        declared_metaclass = builder.add(
+            LoadAddress(type_object_op.type, type_object_op.src, cdef.line)
+        )
+
+    return builder.call_c(py_calc_meta_op, [declared_metaclass, bases], cdef.line)
+
+
+def setup_non_ext_dict(
+    builder: IRBuilder, cdef: ClassDef, metaclass: Value, bases: Value
+) -> Value:
+    """Initialize the class dictionary for a non-extension class.
+
+    This class dictionary is passed to the metaclass constructor.
+    """
+    # Check if the metaclass defines a __prepare__ method, and if so, call it.
+    has_prepare = builder.call_c(
+        py_hasattr_op, [metaclass, builder.load_str("__prepare__")], cdef.line
+    )
+
+    non_ext_dict = Register(dict_rprimitive)
+
+    true_block, false_block, exit_block = BasicBlock(), BasicBlock(), BasicBlock()
+    builder.add_bool_branch(has_prepare, true_block, false_block)
+
+    builder.activate_block(true_block)
+    cls_name = builder.load_str(cdef.name)
+    prepare_meth = builder.py_get_attr(metaclass, "__prepare__", cdef.line)
+    prepare_dict = builder.py_call(prepare_meth, [cls_name, bases], cdef.line)
+    builder.assign(non_ext_dict, prepare_dict, cdef.line)
+    builder.goto(exit_block)
+
+    builder.activate_block(false_block)
+    builder.assign(non_ext_dict, builder.call_c(dict_new_op, [], cdef.line), cdef.line)
+    builder.goto(exit_block)
+    builder.activate_block(exit_block)
+
+    return non_ext_dict
+
+
+def add_non_ext_class_attr_ann(
+    builder: IRBuilder,
+    non_ext: NonExtClassInfo,
+    lvalue: NameExpr,
+    stmt: AssignmentStmt,
+    get_type_info: Callable[[AssignmentStmt], TypeInfo | None] | None = None,
+) -> None:
+    """Add a class attribute to __annotations__ of a non-extension class."""
+    # FIXME: try to better preserve the special forms and type parameters of generics.
+    typ: Value | None = None
+    if get_type_info is not None:
+        type_info = get_type_info(stmt)
+        if type_info:
+            typ = load_type(builder, type_info, stmt.line)
+
+    if typ is None:
+        # FIXME: if get_type_info is not provided, don't fall back to stmt.type?
+        ann_type = get_proper_type(stmt.type)
+        if (
+            isinstance(stmt.unanalyzed_type, UnboundType)
+            and stmt.unanalyzed_type.original_str_expr is not None
+        ):
+            # Annotation is a forward reference, so don't attempt to load the actual
+            # type and load the string instead.
+            #
+            # TODO: is it possible to determine whether a non-string annotation is
+            # actually a forward reference due to the __annotations__ future?
+            typ = builder.load_str(stmt.unanalyzed_type.original_str_expr)
+        elif isinstance(ann_type, Instance):
+            typ = load_type(builder, ann_type.type, stmt.line)
+        else:
+            typ = builder.add(LoadAddress(type_object_op.type, type_object_op.src, stmt.line))
+
+    key = builder.load_str(lvalue.name)
+    builder.call_c(dict_set_item_op, [non_ext.anns, key, typ], stmt.line)
+
+
+def add_non_ext_class_attr(
+    builder: IRBuilder,
+    non_ext: NonExtClassInfo,
+    lvalue: NameExpr,
+    stmt: AssignmentStmt,
+    cdef: ClassDef,
+    attr_to_cache: list[tuple[Lvalue, RType]],
+) -> None:
+    """Add a class attribute to __dict__ of a non-extension class."""
+    # Only add the attribute to the __dict__ if the assignment is of the form:
+    # x: type = value (don't add attributes of the form 'x: type' to the __dict__).
+    if not isinstance(stmt.rvalue, TempNode):
+        rvalue = builder.accept(stmt.rvalue)
+        builder.add_to_non_ext_dict(non_ext, lvalue.name, rvalue, stmt.line)
+        # We cache enum attributes to speed up enum attribute lookup since they
+        # are final.
+        if (
+            cdef.info.bases
+            and cdef.info.bases[0].type.fullname == "enum.Enum"
+            # Skip these since Enum will remove it
+            and lvalue.name not in ENUM_REMOVED_PROPS
+        ):
+            # Enum values are always boxed, so use object_rprimitive.
+            attr_to_cache.append((lvalue, object_rprimitive))
+
+
+def find_attr_initializers(
+    builder: IRBuilder, cdef: ClassDef, skip: Callable[[str, AssignmentStmt], bool] | None = None
+) -> tuple[set[str], list[AssignmentStmt]]:
+    """Find initializers of attributes in a class body.
+
+    If provided, the skip arg should be a callable which will return whether
+    to skip generating a default for an attribute.  It will be passed the name of
+    the attribute and the corresponding AssignmentStmt.
+    """
+    cls = builder.mapper.type_to_ir[cdef.info]
+    if cls.builtin_base:
+        return set(), []
+
+    attrs_with_defaults = set()
+
+    # Pull out all assignments in classes in the mro so we can initialize them
+    # TODO: Support nested statements
+    default_assignments = []
+    for info in reversed(cdef.info.mro):
+        if info not in builder.mapper.type_to_ir:
+            continue
+        for stmt in info.defn.defs.body:
+            if (
+                isinstance(stmt, AssignmentStmt)
+                and isinstance(stmt.lvalues[0], NameExpr)
+                and not is_class_var(stmt.lvalues[0])
+                and not isinstance(stmt.rvalue, TempNode)
+            ):
+                name = stmt.lvalues[0].name
+                if name == "__slots__":
+                    continue
+
+                if name == "__deletable__":
+                    check_deletable_declaration(builder, cls, stmt.line)
+                    continue
+
+                if skip is not None and skip(name, stmt):
+                    continue
+
+                attr_type = cls.attr_type(name)
+
+                # If the attribute is initialized to None and type isn't optional,
+                # doesn't initialize it to anything (special case for "# type:" comments).
+                if isinstance(stmt.rvalue, RefExpr) and stmt.rvalue.fullname == "builtins.None":
+                    if (
+                        not is_optional_type(attr_type)
+                        and not is_object_rprimitive(attr_type)
+                        and not is_none_rprimitive(attr_type)
+                    ):
+                        continue
+
+                attrs_with_defaults.add(name)
+                default_assignments.append(stmt)
+
+    return attrs_with_defaults, default_assignments
+
+
+def generate_attr_defaults_init(
+    builder: IRBuilder, cdef: ClassDef, default_assignments: list[AssignmentStmt]
+) -> None:
+    """Generate an initialization method for default attr values (from class vars)."""
+    if not default_assignments:
+        return
+    cls = builder.mapper.type_to_ir[cdef.info]
+    if cls.builtin_base:
+        return
+
+    with builder.enter_method(cls, "__mypyc_defaults_setup", bool_rprimitive):
+        self_var = builder.self()
+        for stmt in default_assignments:
+            lvalue = stmt.lvalues[0]
+            assert isinstance(lvalue, NameExpr)
+            if not stmt.is_final_def and not is_constant(stmt.rvalue):
+                builder.warning("Unsupported default attribute value", stmt.rvalue.line)
+
+            attr_type = cls.attr_type(lvalue.name)
+            val = builder.coerce(builder.accept(stmt.rvalue), attr_type, stmt.line)
+            init = SetAttr(self_var, lvalue.name, val, -1)
+            init.mark_as_initializer()
+            builder.add(init)
+
+        builder.add(Return(builder.true()))
+
+
+def check_deletable_declaration(builder: IRBuilder, cl: ClassIR, line: int) -> None:
+    for attr in cl.deletable:
+        if attr not in cl.attributes:
+            if not cl.has_attr(attr):
+                builder.error(f'Attribute "{attr}" not defined', line)
+                continue
+            for base in cl.mro:
+                if attr in base.property_types:
+                    builder.error(f'Cannot make property "{attr}" deletable', line)
+                    break
+            else:
+                _, base = cl.attr_details(attr)
+                builder.error(
+                    ('Attribute "{}" not defined in "{}" ' + '(defined in "{}")').format(
+                        attr, cl.name, base.name
+                    ),
+                    line,
+                )
+
+
+def create_ne_from_eq(builder: IRBuilder, cdef: ClassDef) -> None:
+    """Create a "__ne__" method from a "__eq__" method (if only latter exists)."""
+    cls = builder.mapper.type_to_ir[cdef.info]
+    if cls.has_method("__eq__") and not cls.has_method("__ne__"):
+        gen_glue_ne_method(builder, cls, cdef.line)
+
+
+def gen_glue_ne_method(builder: IRBuilder, cls: ClassIR, line: int) -> None:
+    """Generate a "__ne__" method from a "__eq__" method."""
+    with builder.enter_method(cls, "__ne__", object_rprimitive):
+        rhs_arg = builder.add_argument("rhs", object_rprimitive)
+
+        # If __eq__ returns NotImplemented, then __ne__ should also
+        not_implemented_block, regular_block = BasicBlock(), BasicBlock()
+        eqval = builder.add(MethodCall(builder.self(), "__eq__", [rhs_arg], line))
+        not_implemented = builder.add(
+            LoadAddress(not_implemented_op.type, not_implemented_op.src, line)
+        )
+        builder.add(
+            Branch(
+                builder.translate_is_op(eqval, not_implemented, "is", line),
+                not_implemented_block,
+                regular_block,
+                Branch.BOOL,
+            )
+        )
+
+        builder.activate_block(regular_block)
+        retval = builder.coerce(builder.unary_op(eqval, "not", line), object_rprimitive, line)
+        builder.add(Return(retval))
+
+        builder.activate_block(not_implemented_block)
+        builder.add(Return(not_implemented))
+
+
+def load_non_ext_class(
+    builder: IRBuilder, ir: ClassIR, non_ext: NonExtClassInfo, line: int
+) -> Value:
+    cls_name = builder.load_str(ir.name)
+
+    add_dunders_to_non_ext_dict(builder, non_ext, line)
+
+    class_type_obj = builder.py_call(
+        non_ext.metaclass, [cls_name, non_ext.bases, non_ext.dict], line
+    )
+    return class_type_obj
+
+
+def load_decorated_class(builder: IRBuilder, cdef: ClassDef, type_obj: Value) -> Value:
+    """Apply class decorators to create a decorated (non-extension) class object.
+
+    Given a decorated ClassDef and a register containing a
+    non-extension representation of the ClassDef created via the type
+    constructor, applies the corresponding decorator functions on that
+    decorated ClassDef and returns a register containing the decorated
+    ClassDef.
+    """
+    decorators = cdef.decorators
+    dec_class = type_obj
+    for d in reversed(decorators):
+        decorator = d.accept(builder.visitor)
+        assert isinstance(decorator, Value)
+        dec_class = builder.py_call(decorator, [dec_class], dec_class.line)
+    return dec_class
+
+
+def cache_class_attrs(
+    builder: IRBuilder, attrs_to_cache: list[tuple[Lvalue, RType]], cdef: ClassDef
+) -> None:
+    """Add class attributes to be cached to the global cache."""
+    typ = builder.load_native_type_object(cdef.info.fullname)
+    for lval, rtype in attrs_to_cache:
+        assert isinstance(lval, NameExpr)
+        rval = builder.py_get_attr(typ, lval.name, cdef.line)
+        builder.init_final_static(lval, rval, cdef.name, type_override=rtype)
+
+
+def create_mypyc_attrs_tuple(builder: IRBuilder, ir: ClassIR, line: int) -> Value:
+    attrs = [name for ancestor in ir.mro for name in ancestor.attributes]
+    if ir.inherits_python:
+        attrs.append("__dict__")
+    items = [builder.load_str(attr) for attr in attrs]
+    return builder.new_tuple(items, line)
+
+
+def add_dunders_to_non_ext_dict(
+    builder: IRBuilder, non_ext: NonExtClassInfo, line: int, add_annotations: bool = True
+) -> None:
+    if add_annotations:
+        # Add __annotations__ to the class dict.
+        builder.add_to_non_ext_dict(non_ext, "__annotations__", non_ext.anns, line)
+
+    # We add a __doc__ attribute so if the non-extension class is decorated with the
+    # dataclass decorator, dataclass will not try to look for __text_signature__.
+    # https://github.com/python/cpython/blob/3.7/Lib/dataclasses.py#L957
+    filler_doc_str = "mypyc filler docstring"
+    builder.add_to_non_ext_dict(non_ext, "__doc__", builder.load_str(filler_doc_str), line)
+    builder.add_to_non_ext_dict(non_ext, "__module__", builder.load_str(builder.module_name), line)

venv/lib/python3.12/site-packages/mypyc/irbuild/constant_fold.py

@@ -0,0 +1,95 @@
+"""Constant folding of IR values.
+
+For example, 3 + 5 can be constant folded into 8.
+
+This is mostly like mypy.constant_fold, but we can bind some additional
+NameExpr and MemberExpr references here, since we have more knowledge
+about which definitions can be trusted -- we constant fold only references
+to other compiled modules in the same compilation unit.
+"""
+
+from __future__ import annotations
+
+from typing import Final, Union
+
+from mypy.constant_fold import constant_fold_binary_op, constant_fold_unary_op
+from mypy.nodes import (
+    BytesExpr,
+    ComplexExpr,
+    Expression,
+    FloatExpr,
+    IntExpr,
+    MemberExpr,
+    NameExpr,
+    OpExpr,
+    StrExpr,
+    UnaryExpr,
+    Var,
+)
+from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.util import bytes_from_str
+
+# All possible result types of constant folding
+ConstantValue = Union[int, float, complex, str, bytes]
+CONST_TYPES: Final = (int, float, complex, str, bytes)
+
+
+def constant_fold_expr(builder: IRBuilder, expr: Expression) -> ConstantValue | None:
+    """Return the constant value of an expression for supported operations.
+
+    Return None otherwise.
+    """
+    if isinstance(expr, IntExpr):
+        return expr.value
+    if isinstance(expr, FloatExpr):
+        return expr.value
+    if isinstance(expr, StrExpr):
+        return expr.value
+    if isinstance(expr, BytesExpr):
+        return bytes_from_str(expr.value)
+    if isinstance(expr, ComplexExpr):
+        return expr.value
+    elif isinstance(expr, NameExpr):
+        node = expr.node
+        if isinstance(node, Var) and node.is_final:
+            final_value = node.final_value
+            if isinstance(final_value, (CONST_TYPES)):
+                return final_value
+    elif isinstance(expr, MemberExpr):
+        final = builder.get_final_ref(expr)
+        if final is not None:
+            fn, final_var, native = final
+            if final_var.is_final:
+                final_value = final_var.final_value
+                if isinstance(final_value, (CONST_TYPES)):
+                    return final_value
+    elif isinstance(expr, OpExpr):
+        left = constant_fold_expr(builder, expr.left)
+        right = constant_fold_expr(builder, expr.right)
+        if left is not None and right is not None:
+            return constant_fold_binary_op_extended(expr.op, left, right)
+    elif isinstance(expr, UnaryExpr):
+        value = constant_fold_expr(builder, expr.expr)
+        if value is not None and not isinstance(value, bytes):
+            return constant_fold_unary_op(expr.op, value)
+    return None
+
+
+def constant_fold_binary_op_extended(
+    op: str, left: ConstantValue, right: ConstantValue
+) -> ConstantValue | None:
+    """Like mypy's constant_fold_binary_op(), but includes bytes support.
+
+    mypy cannot use constant folded bytes easily so it's simpler to only support them in mypyc.
+    """
+    if not isinstance(left, bytes) and not isinstance(right, bytes):
+        return constant_fold_binary_op(op, left, right)
+
+    if op == "+" and isinstance(left, bytes) and isinstance(right, bytes):
+        return left + right
+    elif op == "*" and isinstance(left, bytes) and isinstance(right, int):
+        return left * right
+    elif op == "*" and isinstance(left, int) and isinstance(right, bytes):
+        return left * right
+
+    return None

venv/lib/python3.12/site-packages/mypyc/irbuild/context.py

@@ -0,0 +1,184 @@
+"""Helpers that store information about functions and the related classes."""
+
+from __future__ import annotations
+
+from mypy.nodes import FuncItem
+from mypyc.ir.class_ir import ClassIR
+from mypyc.ir.func_ir import INVALID_FUNC_DEF
+from mypyc.ir.ops import BasicBlock, Value
+from mypyc.irbuild.targets import AssignmentTarget
+
+
+class FuncInfo:
+    """Contains information about functions as they are generated."""
+
+    def __init__(
+        self,
+        fitem: FuncItem = INVALID_FUNC_DEF,
+        name: str = "",
+        class_name: str | None = None,
+        namespace: str = "",
+        is_nested: bool = False,
+        contains_nested: bool = False,
+        is_decorated: bool = False,
+        in_non_ext: bool = False,
+    ) -> None:
+        self.fitem = fitem
+        self.name = name
+        self.class_name = class_name
+        self.ns = namespace
+        # Callable classes implement the '__call__' method, and are used to represent functions
+        # that are nested inside of other functions.
+        self._callable_class: ImplicitClass | None = None
+        # Environment classes are ClassIR instances that contain attributes representing the
+        # variables in the environment of the function they correspond to. Environment classes are
+        # generated for functions that contain nested functions.
+        self._env_class: ClassIR | None = None
+        # Generator classes implement the '__next__' method, and are used to represent generators
+        # returned by generator functions.
+        self._generator_class: GeneratorClass | None = None
+        # Environment class registers are the local registers associated with instances of an
+        # environment class, used for getting and setting attributes. curr_env_reg is the register
+        # associated with the current environment.
+        self._curr_env_reg: Value | None = None
+        # These are flags denoting whether a given function is nested, contains a nested function,
+        # is decorated, or is within a non-extension class.
+        self.is_nested = is_nested
+        self.contains_nested = contains_nested
+        self.is_decorated = is_decorated
+        self.in_non_ext = in_non_ext
+
+        # TODO: add field for ret_type: RType = none_rprimitive
+
+    def namespaced_name(self) -> str:
+        return "_".join(x for x in [self.name, self.class_name, self.ns] if x)
+
+    @property
+    def is_generator(self) -> bool:
+        return self.fitem.is_generator or self.fitem.is_coroutine
+
+    @property
+    def is_coroutine(self) -> bool:
+        return self.fitem.is_coroutine
+
+    @property
+    def callable_class(self) -> ImplicitClass:
+        assert self._callable_class is not None
+        return self._callable_class
+
+    @callable_class.setter
+    def callable_class(self, cls: ImplicitClass) -> None:
+        self._callable_class = cls
+
+    @property
+    def env_class(self) -> ClassIR:
+        assert self._env_class is not None
+        return self._env_class
+
+    @env_class.setter
+    def env_class(self, ir: ClassIR) -> None:
+        self._env_class = ir
+
+    @property
+    def generator_class(self) -> GeneratorClass:
+        assert self._generator_class is not None
+        return self._generator_class
+
+    @generator_class.setter
+    def generator_class(self, cls: GeneratorClass) -> None:
+        self._generator_class = cls
+
+    @property
+    def curr_env_reg(self) -> Value:
+        assert self._curr_env_reg is not None
+        return self._curr_env_reg
+
+
+class ImplicitClass:
+    """Contains information regarding implicitly generated classes.
+
+    Implicit classes are generated for nested functions and generator
+    functions. They are not explicitly defined in the source code.
+
+    NOTE: This is both a concrete class and used as a base class.
+    """
+
+    def __init__(self, ir: ClassIR) -> None:
+        # The ClassIR instance associated with this class.
+        self.ir = ir
+        # The register associated with the 'self' instance for this generator class.
+        self._self_reg: Value | None = None
+        # Environment class registers are the local registers associated with instances of an
+        # environment class, used for getting and setting attributes. curr_env_reg is the register
+        # associated with the current environment. prev_env_reg is the self.__mypyc_env__ field
+        # associated with the previous environment.
+        self._curr_env_reg: Value | None = None
+        self._prev_env_reg: Value | None = None
+
+    @property
+    def self_reg(self) -> Value:
+        assert self._self_reg is not None
+        return self._self_reg
+
+    @self_reg.setter
+    def self_reg(self, reg: Value) -> None:
+        self._self_reg = reg
+
+    @property
+    def curr_env_reg(self) -> Value:
+        assert self._curr_env_reg is not None
+        return self._curr_env_reg
+
+    @curr_env_reg.setter
+    def curr_env_reg(self, reg: Value) -> None:
+        self._curr_env_reg = reg
+
+    @property
+    def prev_env_reg(self) -> Value:
+        assert self._prev_env_reg is not None
+        return self._prev_env_reg
+
+    @prev_env_reg.setter
+    def prev_env_reg(self, reg: Value) -> None:
+        self._prev_env_reg = reg
+
+
+class GeneratorClass(ImplicitClass):
+    """Contains information about implicit generator function classes."""
+
+    def __init__(self, ir: ClassIR) -> None:
+        super().__init__(ir)
+        # This register holds the label number that the '__next__' function should go to the next
+        # time it is called.
+        self._next_label_reg: Value | None = None
+        self._next_label_target: AssignmentTarget | None = None
+
+        # These registers hold the error values for the generator object for the case that the
+        # 'throw' function is called.
+        self.exc_regs: tuple[Value, Value, Value] | None = None
+
+        # Holds the arg passed to send
+        self.send_arg_reg: Value | None = None
+
+        # The switch block is used to decide which instruction to go using the value held in the
+        # next-label register.
+        self.switch_block = BasicBlock()
+        self.continuation_blocks: list[BasicBlock] = []
+
+    @property
+    def next_label_reg(self) -> Value:
+        assert self._next_label_reg is not None
+        return self._next_label_reg
+
+    @next_label_reg.setter
+    def next_label_reg(self, reg: Value) -> None:
+        self._next_label_reg = reg
+
+    @property
+    def next_label_target(self) -> AssignmentTarget:
+        assert self._next_label_target is not None
+        return self._next_label_target
+
+    @next_label_target.setter
+    def next_label_target(self, target: AssignmentTarget) -> None:
+        self._next_label_target = target

venv/lib/python3.12/site-packages/mypyc/irbuild/env_class.py

@@ -0,0 +1,223 @@
+"""Generate classes representing function environments (+ related operations).
+
+If we have a nested function that has non-local (free) variables, access to the
+non-locals is via an instance of an environment class. Example:
+
+    def f() -> int:
+        x = 0  # Make 'x' an attribute of an environment class instance
+
+        def g() -> int:
+            # We have access to the environment class instance to
+            # allow accessing 'x'
+            return x + 2
+
+        x = x + 1  # Modify the attribute
+        return g()
+"""
+
+from __future__ import annotations
+
+from mypy.nodes import Argument, FuncDef, SymbolNode, Var
+from mypyc.common import BITMAP_BITS, ENV_ATTR_NAME, SELF_NAME, bitmap_name
+from mypyc.ir.class_ir import ClassIR
+from mypyc.ir.ops import Call, GetAttr, SetAttr, Value
+from mypyc.ir.rtypes import RInstance, bitmap_rprimitive, object_rprimitive
+from mypyc.irbuild.builder import IRBuilder, SymbolTarget
+from mypyc.irbuild.context import FuncInfo, GeneratorClass, ImplicitClass
+from mypyc.irbuild.targets import AssignmentTargetAttr
+
+
+def setup_env_class(builder: IRBuilder) -> ClassIR:
+    """Generate a class representing a function environment.
+
+    Note that the variables in the function environment are not
+    actually populated here. This is because when the environment
+    class is generated, the function environment has not yet been
+    visited. This behavior is allowed so that when the compiler visits
+    nested functions, it can use the returned ClassIR instance to
+    figure out free variables it needs to access.  The remaining
+    attributes of the environment class are populated when the
+    environment registers are loaded.
+
+    Return a ClassIR representing an environment for a function
+    containing a nested function.
+    """
+    env_class = ClassIR(
+        f"{builder.fn_info.namespaced_name()}_env", builder.module_name, is_generated=True
+    )
+    env_class.attributes[SELF_NAME] = RInstance(env_class)
+    if builder.fn_info.is_nested:
+        # If the function is nested, its environment class must contain an environment
+        # attribute pointing to its encapsulating functions' environment class.
+        env_class.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_infos[-2].env_class)
+    env_class.mro = [env_class]
+    builder.fn_info.env_class = env_class
+    builder.classes.append(env_class)
+    return env_class
+
+
+def finalize_env_class(builder: IRBuilder) -> None:
+    """Generate, instantiate, and set up the environment of an environment class."""
+    instantiate_env_class(builder)
+
+    # Iterate through the function arguments and replace local definitions (using registers)
+    # that were previously added to the environment with references to the function's
+    # environment class.
+    if builder.fn_info.is_nested:
+        add_args_to_env(builder, local=False, base=builder.fn_info.callable_class)
+    else:
+        add_args_to_env(builder, local=False, base=builder.fn_info)
+
+
+def instantiate_env_class(builder: IRBuilder) -> Value:
+    """Assign an environment class to a register named after the given function definition."""
+    curr_env_reg = builder.add(
+        Call(builder.fn_info.env_class.ctor, [], builder.fn_info.fitem.line)
+    )
+
+    if builder.fn_info.is_nested:
+        builder.fn_info.callable_class._curr_env_reg = curr_env_reg
+        builder.add(
+            SetAttr(
+                curr_env_reg,
+                ENV_ATTR_NAME,
+                builder.fn_info.callable_class.prev_env_reg,
+                builder.fn_info.fitem.line,
+            )
+        )
+    else:
+        builder.fn_info._curr_env_reg = curr_env_reg
+
+    return curr_env_reg
+
+
+def load_env_registers(builder: IRBuilder) -> None:
+    """Load the registers for the current FuncItem being visited.
+
+    Adds the arguments of the FuncItem to the environment. If the
+    FuncItem is nested inside of another function, then this also
+    loads all of the outer environments of the FuncItem into registers
+    so that they can be used when accessing free variables.
+    """
+    add_args_to_env(builder, local=True)
+
+    fn_info = builder.fn_info
+    fitem = fn_info.fitem
+    if fn_info.is_nested:
+        load_outer_envs(builder, fn_info.callable_class)
+        # If this is a FuncDef, then make sure to load the FuncDef into its own environment
+        # class so that the function can be called recursively.
+        if isinstance(fitem, FuncDef):
+            setup_func_for_recursive_call(builder, fitem, fn_info.callable_class)
+
+
+def load_outer_env(
+    builder: IRBuilder, base: Value, outer_env: dict[SymbolNode, SymbolTarget]
+) -> Value:
+    """Load the environment class for a given base into a register.
+
+    Additionally, iterates through all of the SymbolNode and
+    AssignmentTarget instances of the environment at the given index's
+    symtable, and adds those instances to the environment of the
+    current environment. This is done so that the current environment
+    can access outer environment variables without having to reload
+    all of the environment registers.
+
+    Returns the register where the environment class was loaded.
+    """
+    env = builder.add(GetAttr(base, ENV_ATTR_NAME, builder.fn_info.fitem.line))
+    assert isinstance(env.type, RInstance), f"{env} must be of type RInstance"
+
+    for symbol, target in outer_env.items():
+        env.type.class_ir.attributes[symbol.name] = target.type
+        symbol_target = AssignmentTargetAttr(env, symbol.name)
+        builder.add_target(symbol, symbol_target)
+
+    return env
+
+
+def load_outer_envs(builder: IRBuilder, base: ImplicitClass) -> None:
+    index = len(builder.builders) - 2
+
+    # Load the first outer environment. This one is special because it gets saved in the
+    # FuncInfo instance's prev_env_reg field.
+    if index > 1:
+        # outer_env = builder.fn_infos[index].environment
+        outer_env = builder.symtables[index]
+        if isinstance(base, GeneratorClass):
+            base.prev_env_reg = load_outer_env(builder, base.curr_env_reg, outer_env)
+        else:
+            base.prev_env_reg = load_outer_env(builder, base.self_reg, outer_env)
+        env_reg = base.prev_env_reg
+        index -= 1
+
+    # Load the remaining outer environments into registers.
+    while index > 1:
+        # outer_env = builder.fn_infos[index].environment
+        outer_env = builder.symtables[index]
+        env_reg = load_outer_env(builder, env_reg, outer_env)
+        index -= 1
+
+
+def num_bitmap_args(builder: IRBuilder, args: list[Argument]) -> int:
+    n = 0
+    for arg in args:
+        t = builder.type_to_rtype(arg.variable.type)
+        if t.error_overlap and arg.kind.is_optional():
+            n += 1
+    return (n + (BITMAP_BITS - 1)) // BITMAP_BITS
+
+
+def add_args_to_env(
+    builder: IRBuilder,
+    local: bool = True,
+    base: FuncInfo | ImplicitClass | None = None,
+    reassign: bool = True,
+) -> None:
+    fn_info = builder.fn_info
+    args = fn_info.fitem.arguments
+    nb = num_bitmap_args(builder, args)
+    if local:
+        for arg in args:
+            rtype = builder.type_to_rtype(arg.variable.type)
+            builder.add_local_reg(arg.variable, rtype, is_arg=True)
+        for i in reversed(range(nb)):
+            builder.add_local_reg(Var(bitmap_name(i)), bitmap_rprimitive, is_arg=True)
+    else:
+        for arg in args:
+            if is_free_variable(builder, arg.variable) or fn_info.is_generator:
+                rtype = builder.type_to_rtype(arg.variable.type)
+                assert base is not None, "base cannot be None for adding nonlocal args"
+                builder.add_var_to_env_class(arg.variable, rtype, base, reassign=reassign)
+
+
+def setup_func_for_recursive_call(builder: IRBuilder, fdef: FuncDef, base: ImplicitClass) -> None:
+    """Enable calling a nested function (with a callable class) recursively.
+
+    Adds the instance of the callable class representing the given
+    FuncDef to a register in the environment so that the function can
+    be called recursively. Note that this needs to be done only for
+    nested functions.
+    """
+    # First, set the attribute of the environment class so that GetAttr can be called on it.
+    prev_env = builder.fn_infos[-2].env_class
+    prev_env.attributes[fdef.name] = builder.type_to_rtype(fdef.type)
+
+    if isinstance(base, GeneratorClass):
+        # If we are dealing with a generator class, then we need to first get the register
+        # holding the current environment class, and load the previous environment class from
+        # there.
+        prev_env_reg = builder.add(GetAttr(base.curr_env_reg, ENV_ATTR_NAME, -1))
+    else:
+        prev_env_reg = base.prev_env_reg
+
+    # Obtain the instance of the callable class representing the FuncDef, and add it to the
+    # current environment.
+    val = builder.add(GetAttr(prev_env_reg, fdef.name, -1))
+    target = builder.add_local_reg(fdef, object_rprimitive)
+    builder.assign(target, val, -1)
+
+
+def is_free_variable(builder: IRBuilder, symbol: SymbolNode) -> bool:
+    fitem = builder.fn_info.fitem
+    return fitem in builder.free_variables and symbol in builder.free_variables[fitem]

venv/lib/python3.12/site-packages/mypyc/irbuild/format_str_tokenizer.py

@@ -0,0 +1,250 @@
+"""Tokenizers for three string formatting methods"""
+
+from __future__ import annotations
+
+from enum import Enum, unique
+from typing import Final
+
+from mypy.checkstrformat import (
+    ConversionSpecifier,
+    parse_conversion_specifiers,
+    parse_format_value,
+)
+from mypy.errors import Errors
+from mypy.messages import MessageBuilder
+from mypy.nodes import Context, Expression
+from mypy.options import Options
+from mypyc.ir.ops import Integer, Value
+from mypyc.ir.rtypes import (
+    c_pyssize_t_rprimitive,
+    is_bytes_rprimitive,
+    is_int_rprimitive,
+    is_short_int_rprimitive,
+    is_str_rprimitive,
+)
+from mypyc.irbuild.builder import IRBuilder
+from mypyc.primitives.bytes_ops import bytes_build_op
+from mypyc.primitives.int_ops import int_to_str_op
+from mypyc.primitives.str_ops import str_build_op, str_op
+
+
+@unique
+class FormatOp(Enum):
+    """FormatOp represents conversion operations of string formatting during
+    compile time.
+
+    Compare to ConversionSpecifier, FormatOp has fewer attributes.
+    For example, to mark a conversion from any object to string,
+    ConversionSpecifier may have several representations, like '%s', '{}'
+    or '{:{}}'. However, there would only exist one corresponding FormatOp.
+    """
+
+    STR = "s"
+    INT = "d"
+    BYTES = "b"
+
+
+def generate_format_ops(specifiers: list[ConversionSpecifier]) -> list[FormatOp] | None:
+    """Convert ConversionSpecifier to FormatOp.
+
+    Different ConversionSpecifiers may share a same FormatOp.
+    """
+    format_ops = []
+    for spec in specifiers:
+        # TODO: Match specifiers instead of using whole_seq
+        if spec.whole_seq == "%s" or spec.whole_seq == "{:{}}":
+            format_op = FormatOp.STR
+        elif spec.whole_seq == "%d":
+            format_op = FormatOp.INT
+        elif spec.whole_seq == "%b":
+            format_op = FormatOp.BYTES
+        elif spec.whole_seq:
+            return None
+        else:
+            format_op = FormatOp.STR
+        format_ops.append(format_op)
+    return format_ops
+
+
+def tokenizer_printf_style(format_str: str) -> tuple[list[str], list[FormatOp]] | None:
+    """Tokenize a printf-style format string using regex.
+
+    Return:
+        A list of string literals and a list of FormatOps.
+    """
+    literals: list[str] = []
+    specifiers: list[ConversionSpecifier] = parse_conversion_specifiers(format_str)
+    format_ops = generate_format_ops(specifiers)
+    if format_ops is None:
+        return None
+
+    last_end = 0
+    for spec in specifiers:
+        cur_start = spec.start_pos
+        literals.append(format_str[last_end:cur_start])
+        last_end = cur_start + len(spec.whole_seq)
+    literals.append(format_str[last_end:])
+
+    return literals, format_ops
+
+
+# The empty Context as an argument for parse_format_value().
+# It wouldn't be used since the code has passed the type-checking.
+EMPTY_CONTEXT: Final = Context()
+
+
+def tokenizer_format_call(format_str: str) -> tuple[list[str], list[FormatOp]] | None:
+    """Tokenize a str.format() format string.
+
+    The core function parse_format_value() is shared with mypy.
+    With these specifiers, we then parse the literal substrings
+    of the original format string and convert `ConversionSpecifier`
+    to `FormatOp`.
+
+    Return:
+        A list of string literals and a list of FormatOps. The literals
+        are interleaved with FormatOps and the length of returned literals
+        should be exactly one more than FormatOps.
+        Return None if it cannot parse the string.
+    """
+    # Creates an empty MessageBuilder here.
+    # It wouldn't be used since the code has passed the type-checking.
+    specifiers = parse_format_value(
+        format_str, EMPTY_CONTEXT, MessageBuilder(Errors(Options()), {})
+    )
+    if specifiers is None:
+        return None
+    format_ops = generate_format_ops(specifiers)
+    if format_ops is None:
+        return None
+
+    literals: list[str] = []
+    last_end = 0
+    for spec in specifiers:
+        # Skip { and }
+        literals.append(format_str[last_end : spec.start_pos - 1])
+        last_end = spec.start_pos + len(spec.whole_seq) + 1
+    literals.append(format_str[last_end:])
+    # Deal with escaped {{
+    literals = [x.replace("{{", "{").replace("}}", "}") for x in literals]
+
+    return literals, format_ops
+
+
+def convert_format_expr_to_str(
+    builder: IRBuilder, format_ops: list[FormatOp], exprs: list[Expression], line: int
+) -> list[Value] | None:
+    """Convert expressions into string literal objects with the guidance
+    of FormatOps. Return None when fails."""
+    if len(format_ops) != len(exprs):
+        return None
+
+    converted = []
+    for x, format_op in zip(exprs, format_ops):
+        node_type = builder.node_type(x)
+        if format_op == FormatOp.STR:
+            if is_str_rprimitive(node_type):
+                var_str = builder.accept(x)
+            elif is_int_rprimitive(node_type) or is_short_int_rprimitive(node_type):
+                var_str = builder.call_c(int_to_str_op, [builder.accept(x)], line)
+            else:
+                var_str = builder.call_c(str_op, [builder.accept(x)], line)
+        elif format_op == FormatOp.INT:
+            if is_int_rprimitive(node_type) or is_short_int_rprimitive(node_type):
+                var_str = builder.call_c(int_to_str_op, [builder.accept(x)], line)
+            else:
+                return None
+        else:
+            return None
+        converted.append(var_str)
+    return converted
+
+
+def join_formatted_strings(
+    builder: IRBuilder, literals: list[str] | None, substitutions: list[Value], line: int
+) -> Value:
+    """Merge the list of literals and the list of substitutions
+    alternatively using 'str_build_op'.
+
+    `substitutions` is the result value of formatting conversions.
+
+    If the `literals` is set to None, we simply join the substitutions;
+    Otherwise, the `literals` is the literal substrings of the original
+    format string and its length should be exactly one more than
+    substitutions.
+
+    For example:
+    (1)    'This is a %s and the value is %d'
+        -> literals: ['This is a ', ' and the value is', '']
+    (2)    '{} and the value is {}'
+        -> literals: ['', ' and the value is', '']
+    """
+    # The first parameter for str_build_op is the total size of
+    # the following PyObject*
+    result_list: list[Value] = [Integer(0, c_pyssize_t_rprimitive)]
+
+    if literals is not None:
+        for a, b in zip(literals, substitutions):
+            if a:
+                result_list.append(builder.load_str(a))
+            result_list.append(b)
+        if literals[-1]:
+            result_list.append(builder.load_str(literals[-1]))
+    else:
+        result_list.extend(substitutions)
+
+    # Special case for empty string and literal string
+    if len(result_list) == 1:
+        return builder.load_str("")
+    if not substitutions and len(result_list) == 2:
+        return result_list[1]
+
+    result_list[0] = Integer(len(result_list) - 1, c_pyssize_t_rprimitive)
+    return builder.call_c(str_build_op, result_list, line)
+
+
+def convert_format_expr_to_bytes(
+    builder: IRBuilder, format_ops: list[FormatOp], exprs: list[Expression], line: int
+) -> list[Value] | None:
+    """Convert expressions into bytes literal objects with the guidance
+    of FormatOps. Return None when fails."""
+    if len(format_ops) != len(exprs):
+        return None
+
+    converted = []
+    for x, format_op in zip(exprs, format_ops):
+        node_type = builder.node_type(x)
+        # conversion type 's' is an alias of 'b' in bytes formatting
+        if format_op == FormatOp.BYTES or format_op == FormatOp.STR:
+            if is_bytes_rprimitive(node_type):
+                var_bytes = builder.accept(x)
+            else:
+                return None
+        else:
+            return None
+        converted.append(var_bytes)
+    return converted
+
+
+def join_formatted_bytes(
+    builder: IRBuilder, literals: list[str], substitutions: list[Value], line: int
+) -> Value:
+    """Merge the list of literals and the list of substitutions
+    alternatively using 'bytes_build_op'."""
+    result_list: list[Value] = [Integer(0, c_pyssize_t_rprimitive)]
+
+    for a, b in zip(literals, substitutions):
+        if a:
+            result_list.append(builder.load_bytes_from_str_literal(a))
+        result_list.append(b)
+    if literals[-1]:
+        result_list.append(builder.load_bytes_from_str_literal(literals[-1]))
+
+    # Special case for empty bytes and literal
+    if len(result_list) == 1:
+        return builder.load_bytes_from_str_literal("")
+    if not substitutions and len(result_list) == 2:
+        return result_list[1]
+
+    result_list[0] = Integer(len(result_list) - 1, c_pyssize_t_rprimitive)
+    return builder.call_c(bytes_build_op, result_list, line)

venv/lib/python3.12/site-packages/mypyc/irbuild/generator.py

@@ -0,0 +1,346 @@
+"""Generate IR for generator functions.
+
+A generator function is represented by a class that implements the
+generator protocol and keeps track of the generator state, including
+local variables.
+
+The top-level logic for dealing with generator functions is in
+mypyc.irbuild.function.
+"""
+
+from __future__ import annotations
+
+from mypy.nodes import ARG_OPT, Var
+from mypyc.common import ENV_ATTR_NAME, NEXT_LABEL_ATTR_NAME, SELF_NAME
+from mypyc.ir.class_ir import ClassIR
+from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature, RuntimeArg
+from mypyc.ir.ops import (
+    NO_TRACEBACK_LINE_NO,
+    BasicBlock,
+    Branch,
+    Call,
+    Goto,
+    Integer,
+    MethodCall,
+    RaiseStandardError,
+    Register,
+    Return,
+    SetAttr,
+    TupleSet,
+    Unreachable,
+    Value,
+)
+from mypyc.ir.rtypes import RInstance, int_rprimitive, object_rprimitive
+from mypyc.irbuild.builder import IRBuilder, gen_arg_defaults
+from mypyc.irbuild.context import FuncInfo, GeneratorClass
+from mypyc.irbuild.env_class import (
+    add_args_to_env,
+    finalize_env_class,
+    load_env_registers,
+    load_outer_env,
+)
+from mypyc.irbuild.nonlocalcontrol import ExceptNonlocalControl
+from mypyc.primitives.exc_ops import (
+    error_catch_op,
+    exc_matches_op,
+    raise_exception_with_tb_op,
+    reraise_exception_op,
+    restore_exc_info_op,
+)
+
+
+def gen_generator_func(builder: IRBuilder) -> None:
+    setup_generator_class(builder)
+    load_env_registers(builder)
+    gen_arg_defaults(builder)
+    finalize_env_class(builder)
+    builder.add(Return(instantiate_generator_class(builder)))
+
+
+def instantiate_generator_class(builder: IRBuilder) -> Value:
+    fitem = builder.fn_info.fitem
+    generator_reg = builder.add(Call(builder.fn_info.generator_class.ir.ctor, [], fitem.line))
+
+    # Get the current environment register. If the current function is nested, then the
+    # generator class gets instantiated from the callable class' '__call__' method, and hence
+    # we use the callable class' environment register. Otherwise, we use the original
+    # function's environment register.
+    if builder.fn_info.is_nested:
+        curr_env_reg = builder.fn_info.callable_class.curr_env_reg
+    else:
+        curr_env_reg = builder.fn_info.curr_env_reg
+
+    # Set the generator class' environment attribute to point at the environment class
+    # defined in the current scope.
+    builder.add(SetAttr(generator_reg, ENV_ATTR_NAME, curr_env_reg, fitem.line))
+
+    # Set the generator class' environment class' NEXT_LABEL_ATTR_NAME attribute to 0.
+    zero = Integer(0)
+    builder.add(SetAttr(curr_env_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line))
+    return generator_reg
+
+
+def setup_generator_class(builder: IRBuilder) -> ClassIR:
+    name = f"{builder.fn_info.namespaced_name()}_gen"
+
+    generator_class_ir = ClassIR(name, builder.module_name, is_generated=True)
+    generator_class_ir.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_info.env_class)
+    generator_class_ir.mro = [generator_class_ir]
+
+    builder.classes.append(generator_class_ir)
+    builder.fn_info.generator_class = GeneratorClass(generator_class_ir)
+    return generator_class_ir
+
+
+def create_switch_for_generator_class(builder: IRBuilder) -> None:
+    builder.add(Goto(builder.fn_info.generator_class.switch_block))
+    block = BasicBlock()
+    builder.fn_info.generator_class.continuation_blocks.append(block)
+    builder.activate_block(block)
+
+
+def populate_switch_for_generator_class(builder: IRBuilder) -> None:
+    cls = builder.fn_info.generator_class
+    line = builder.fn_info.fitem.line
+
+    builder.activate_block(cls.switch_block)
+    for label, true_block in enumerate(cls.continuation_blocks):
+        false_block = BasicBlock()
+        comparison = builder.binary_op(cls.next_label_reg, Integer(label), "==", line)
+        builder.add_bool_branch(comparison, true_block, false_block)
+        builder.activate_block(false_block)
+
+    builder.add(RaiseStandardError(RaiseStandardError.STOP_ITERATION, None, line))
+    builder.add(Unreachable())
+
+
+def add_raise_exception_blocks_to_generator_class(builder: IRBuilder, line: int) -> None:
+    """Add error handling blocks to a generator class.
+
+    Generates blocks to check if error flags are set while calling the
+    helper method for generator functions, and raises an exception if
+    those flags are set.
+    """
+    cls = builder.fn_info.generator_class
+    assert cls.exc_regs is not None
+    exc_type, exc_val, exc_tb = cls.exc_regs
+
+    # Check to see if an exception was raised.
+    error_block = BasicBlock()
+    ok_block = BasicBlock()
+    comparison = builder.translate_is_op(exc_type, builder.none_object(), "is not", line)
+    builder.add_bool_branch(comparison, error_block, ok_block)
+
+    builder.activate_block(error_block)
+    builder.call_c(raise_exception_with_tb_op, [exc_type, exc_val, exc_tb], line)
+    builder.add(Unreachable())
+    builder.goto_and_activate(ok_block)
+
+
+def add_methods_to_generator_class(
+    builder: IRBuilder,
+    fn_info: FuncInfo,
+    sig: FuncSignature,
+    arg_regs: list[Register],
+    blocks: list[BasicBlock],
+    is_coroutine: bool,
+) -> None:
+    helper_fn_decl = add_helper_to_generator_class(builder, arg_regs, blocks, sig, fn_info)
+    add_next_to_generator_class(builder, fn_info, helper_fn_decl, sig)
+    add_send_to_generator_class(builder, fn_info, helper_fn_decl, sig)
+    add_iter_to_generator_class(builder, fn_info)
+    add_throw_to_generator_class(builder, fn_info, helper_fn_decl, sig)
+    add_close_to_generator_class(builder, fn_info)
+    if is_coroutine:
+        add_await_to_generator_class(builder, fn_info)
+
+
+def add_helper_to_generator_class(
+    builder: IRBuilder,
+    arg_regs: list[Register],
+    blocks: list[BasicBlock],
+    sig: FuncSignature,
+    fn_info: FuncInfo,
+) -> FuncDecl:
+    """Generates a helper method for a generator class, called by '__next__' and 'throw'."""
+    sig = FuncSignature(
+        (
+            RuntimeArg(SELF_NAME, object_rprimitive),
+            RuntimeArg("type", object_rprimitive),
+            RuntimeArg("value", object_rprimitive),
+            RuntimeArg("traceback", object_rprimitive),
+            RuntimeArg("arg", object_rprimitive),
+        ),
+        sig.ret_type,
+    )
+    helper_fn_decl = FuncDecl(
+        "__mypyc_generator_helper__", fn_info.generator_class.ir.name, builder.module_name, sig
+    )
+    helper_fn_ir = FuncIR(
+        helper_fn_decl, arg_regs, blocks, fn_info.fitem.line, traceback_name=fn_info.fitem.name
+    )
+    fn_info.generator_class.ir.methods["__mypyc_generator_helper__"] = helper_fn_ir
+    builder.functions.append(helper_fn_ir)
+    return helper_fn_decl
+
+
+def add_iter_to_generator_class(builder: IRBuilder, fn_info: FuncInfo) -> None:
+    """Generates the '__iter__' method for a generator class."""
+    with builder.enter_method(fn_info.generator_class.ir, "__iter__", object_rprimitive, fn_info):
+        builder.add(Return(builder.self()))
+
+
+def add_next_to_generator_class(
+    builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl, sig: FuncSignature
+) -> None:
+    """Generates the '__next__' method for a generator class."""
+    with builder.enter_method(fn_info.generator_class.ir, "__next__", object_rprimitive, fn_info):
+        none_reg = builder.none_object()
+        # Call the helper function with error flags set to Py_None, and return that result.
+        result = builder.add(
+            Call(
+                fn_decl,
+                [builder.self(), none_reg, none_reg, none_reg, none_reg],
+                fn_info.fitem.line,
+            )
+        )
+        builder.add(Return(result))
+
+
+def add_send_to_generator_class(
+    builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl, sig: FuncSignature
+) -> None:
+    """Generates the 'send' method for a generator class."""
+    with builder.enter_method(fn_info.generator_class.ir, "send", object_rprimitive, fn_info):
+        arg = builder.add_argument("arg", object_rprimitive)
+        none_reg = builder.none_object()
+        # Call the helper function with error flags set to Py_None, and return that result.
+        result = builder.add(
+            Call(
+                fn_decl,
+                [builder.self(), none_reg, none_reg, none_reg, builder.read(arg)],
+                fn_info.fitem.line,
+            )
+        )
+        builder.add(Return(result))
+
+
+def add_throw_to_generator_class(
+    builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl, sig: FuncSignature
+) -> None:
+    """Generates the 'throw' method for a generator class."""
+    with builder.enter_method(fn_info.generator_class.ir, "throw", object_rprimitive, fn_info):
+        typ = builder.add_argument("type", object_rprimitive)
+        val = builder.add_argument("value", object_rprimitive, ARG_OPT)
+        tb = builder.add_argument("traceback", object_rprimitive, ARG_OPT)
+
+        # Because the value and traceback arguments are optional and hence
+        # can be NULL if not passed in, we have to assign them Py_None if
+        # they are not passed in.
+        none_reg = builder.none_object()
+        builder.assign_if_null(val, lambda: none_reg, builder.fn_info.fitem.line)
+        builder.assign_if_null(tb, lambda: none_reg, builder.fn_info.fitem.line)
+
+        # Call the helper function using the arguments passed in, and return that result.
+        result = builder.add(
+            Call(
+                fn_decl,
+                [builder.self(), builder.read(typ), builder.read(val), builder.read(tb), none_reg],
+                fn_info.fitem.line,
+            )
+        )
+        builder.add(Return(result))
+
+
+def add_close_to_generator_class(builder: IRBuilder, fn_info: FuncInfo) -> None:
+    """Generates the '__close__' method for a generator class."""
+    with builder.enter_method(fn_info.generator_class.ir, "close", object_rprimitive, fn_info):
+        except_block, else_block = BasicBlock(), BasicBlock()
+        builder.builder.push_error_handler(except_block)
+        builder.goto_and_activate(BasicBlock())
+        generator_exit = builder.load_module_attr_by_fullname(
+            "builtins.GeneratorExit", fn_info.fitem.line
+        )
+        builder.add(
+            MethodCall(
+                builder.self(),
+                "throw",
+                [generator_exit, builder.none_object(), builder.none_object()],
+            )
+        )
+        builder.goto(else_block)
+        builder.builder.pop_error_handler()
+
+        builder.activate_block(except_block)
+        old_exc = builder.call_c(error_catch_op, [], fn_info.fitem.line)
+        builder.nonlocal_control.append(
+            ExceptNonlocalControl(builder.nonlocal_control[-1], old_exc)
+        )
+        stop_iteration = builder.load_module_attr_by_fullname(
+            "builtins.StopIteration", fn_info.fitem.line
+        )
+        exceptions = builder.add(TupleSet([generator_exit, stop_iteration], fn_info.fitem.line))
+        matches = builder.call_c(exc_matches_op, [exceptions], fn_info.fitem.line)
+
+        match_block, non_match_block = BasicBlock(), BasicBlock()
+        builder.add(Branch(matches, match_block, non_match_block, Branch.BOOL))
+
+        builder.activate_block(match_block)
+        builder.call_c(restore_exc_info_op, [builder.read(old_exc)], fn_info.fitem.line)
+        builder.add(Return(builder.none_object()))
+
+        builder.activate_block(non_match_block)
+        builder.call_c(reraise_exception_op, [], NO_TRACEBACK_LINE_NO)
+        builder.add(Unreachable())
+
+        builder.nonlocal_control.pop()
+
+        builder.activate_block(else_block)
+        builder.add(
+            RaiseStandardError(
+                RaiseStandardError.RUNTIME_ERROR,
+                "generator ignored GeneratorExit",
+                fn_info.fitem.line,
+            )
+        )
+        builder.add(Unreachable())
+
+
+def add_await_to_generator_class(builder: IRBuilder, fn_info: FuncInfo) -> None:
+    """Generates the '__await__' method for a generator class."""
+    with builder.enter_method(fn_info.generator_class.ir, "__await__", object_rprimitive, fn_info):
+        builder.add(Return(builder.self()))
+
+
+def setup_env_for_generator_class(builder: IRBuilder) -> None:
+    """Populates the environment for a generator class."""
+    fitem = builder.fn_info.fitem
+    cls = builder.fn_info.generator_class
+    self_target = builder.add_self_to_env(cls.ir)
+
+    # Add the type, value, and traceback variables to the environment.
+    exc_type = builder.add_local(Var("type"), object_rprimitive, is_arg=True)
+    exc_val = builder.add_local(Var("value"), object_rprimitive, is_arg=True)
+    exc_tb = builder.add_local(Var("traceback"), object_rprimitive, is_arg=True)
+    # TODO: Use the right type here instead of object?
+    exc_arg = builder.add_local(Var("arg"), object_rprimitive, is_arg=True)
+
+    cls.exc_regs = (exc_type, exc_val, exc_tb)
+    cls.send_arg_reg = exc_arg
+
+    cls.self_reg = builder.read(self_target, fitem.line)
+    cls.curr_env_reg = load_outer_env(builder, cls.self_reg, builder.symtables[-1])
+
+    # Define a variable representing the label to go to the next time
+    # the '__next__' function of the generator is called, and add it
+    # as an attribute to the environment class.
+    cls.next_label_target = builder.add_var_to_env_class(
+        Var(NEXT_LABEL_ATTR_NAME), int_rprimitive, cls, reassign=False
+    )
+
+    # Add arguments from the original generator function to the
+    # environment of the generator class.
+    add_args_to_env(builder, local=False, base=cls, reassign=False)
+
+    # Set the next label register for the generator class.
+    cls.next_label_reg = builder.read(cls.next_label_target, fitem.line)

venv/lib/python3.12/site-packages/mypyc/irbuild/main.py

@@ -0,0 +1,153 @@
+"""Transform a mypy AST to the IR form (Intermediate Representation).
+
+For example, consider a function like this:
+
+   def f(x: int) -> int:
+       return x * 2 + 1
+
+It would be translated to something that conceptually looks like this:
+
+   r0 = 2
+   r1 = 1
+   r2 = x * r0 :: int
+   r3 = r2 + r1 :: int
+   return r3
+
+This module deals with the module-level IR transformation logic and
+putting it all together. The actual IR is implemented in mypyc.ir.
+
+For the core of the IR transform implementation, look at build_ir()
+below, mypyc.irbuild.builder, and mypyc.irbuild.visitor.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, TypeVar, cast
+
+from mypy.build import Graph
+from mypy.nodes import ClassDef, Expression, MypyFile
+from mypy.state import state
+from mypy.types import Type
+from mypyc.analysis.attrdefined import analyze_always_defined_attrs
+from mypyc.common import TOP_LEVEL_NAME
+from mypyc.errors import Errors
+from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature
+from mypyc.ir.module_ir import ModuleIR, ModuleIRs
+from mypyc.ir.rtypes import none_rprimitive
+from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.mapper import Mapper
+from mypyc.irbuild.prebuildvisitor import PreBuildVisitor
+from mypyc.irbuild.prepare import build_type_map, find_singledispatch_register_impls
+from mypyc.irbuild.visitor import IRBuilderVisitor
+from mypyc.irbuild.vtable import compute_vtable
+from mypyc.options import CompilerOptions
+
+# The stubs for callable contextmanagers are busted so cast it to the
+# right type...
+F = TypeVar("F", bound=Callable[..., Any])
+strict_optional_dec = cast(Callable[[F], F], state.strict_optional_set(True))
+
+
+@strict_optional_dec  # Turn on strict optional for any type manipulations we do
+def build_ir(
+    modules: list[MypyFile],
+    graph: Graph,
+    types: dict[Expression, Type],
+    mapper: Mapper,
+    options: CompilerOptions,
+    errors: Errors,
+) -> ModuleIRs:
+    """Build basic IR for a set of modules that have been type-checked by mypy.
+
+    The returned IR is not complete and requires additional
+    transformations, such as the insertion of refcount handling.
+    """
+
+    build_type_map(mapper, modules, graph, types, options, errors)
+    singledispatch_info = find_singledispatch_register_impls(modules, errors)
+
+    result: ModuleIRs = {}
+
+    # Generate IR for all modules.
+    class_irs = []
+
+    for module in modules:
+        # First pass to determine free symbols.
+        pbv = PreBuildVisitor(errors, module, singledispatch_info.decorators_to_remove)
+        module.accept(pbv)
+
+        # Construct and configure builder objects (cyclic runtime dependency).
+        visitor = IRBuilderVisitor()
+        builder = IRBuilder(
+            module.fullname,
+            types,
+            graph,
+            errors,
+            mapper,
+            pbv,
+            visitor,
+            options,
+            singledispatch_info.singledispatch_impls,
+        )
+        visitor.builder = builder
+
+        # Second pass does the bulk of the work.
+        transform_mypy_file(builder, module)
+        module_ir = ModuleIR(
+            module.fullname,
+            list(builder.imports),
+            builder.functions,
+            builder.classes,
+            builder.final_names,
+        )
+        result[module.fullname] = module_ir
+        class_irs.extend(builder.classes)
+
+    analyze_always_defined_attrs(class_irs)
+
+    # Compute vtables.
+    for cir in class_irs:
+        if cir.is_ext_class:
+            compute_vtable(cir)
+
+    return result
+
+
+def transform_mypy_file(builder: IRBuilder, mypyfile: MypyFile) -> None:
+    """Generate IR for a single module."""
+
+    if mypyfile.fullname in ("typing", "abc"):
+        # These module are special; their contents are currently all
+        # built-in primitives.
+        return
+
+    builder.set_module(mypyfile.fullname, mypyfile.path)
+
+    classes = [node for node in mypyfile.defs if isinstance(node, ClassDef)]
+
+    # Collect all classes.
+    for cls in classes:
+        ir = builder.mapper.type_to_ir[cls.info]
+        builder.classes.append(ir)
+
+    builder.enter("<module>")
+
+    # Make sure we have a builtins import
+    builder.gen_import("builtins", -1)
+
+    # Generate ops.
+    for node in mypyfile.defs:
+        builder.accept(node)
+
+    builder.maybe_add_implicit_return()
+
+    # Generate special function representing module top level.
+    args, _, blocks, ret_type, _ = builder.leave()
+    sig = FuncSignature([], none_rprimitive)
+    func_ir = FuncIR(
+        FuncDecl(TOP_LEVEL_NAME, None, builder.module_name, sig),
+        args,
+        blocks,
+        traceback_name="<module>",
+    )
+    builder.functions.append(func_ir)

venv/lib/python3.12/site-packages/mypyc/irbuild/mapper.py

@@ -0,0 +1,217 @@
+"""Maintain a mapping from mypy concepts to IR/compiled concepts."""
+
+from __future__ import annotations
+
+from mypy.nodes import ARG_STAR, ARG_STAR2, GDEF, ArgKind, FuncDef, RefExpr, SymbolNode, TypeInfo
+from mypy.types import (
+    AnyType,
+    CallableType,
+    Instance,
+    LiteralType,
+    NoneTyp,
+    Overloaded,
+    PartialType,
+    TupleType,
+    Type,
+    TypedDictType,
+    TypeType,
+    TypeVarType,
+    UnboundType,
+    UninhabitedType,
+    UnionType,
+    get_proper_type,
+)
+from mypyc.ir.class_ir import ClassIR
+from mypyc.ir.func_ir import FuncDecl, FuncSignature, RuntimeArg
+from mypyc.ir.rtypes import (
+    RInstance,
+    RTuple,
+    RType,
+    RUnion,
+    bool_rprimitive,
+    bytes_rprimitive,
+    dict_rprimitive,
+    float_rprimitive,
+    int16_rprimitive,
+    int32_rprimitive,
+    int64_rprimitive,
+    int_rprimitive,
+    list_rprimitive,
+    none_rprimitive,
+    object_rprimitive,
+    range_rprimitive,
+    set_rprimitive,
+    str_rprimitive,
+    tuple_rprimitive,
+    uint8_rprimitive,
+)
+
+
+class Mapper:
+    """Keep track of mappings from mypy concepts to IR concepts.
+
+    For example, we keep track of how the mypy TypeInfos of compiled
+    classes map to class IR objects.
+
+    This state is shared across all modules being compiled in all
+    compilation groups.
+    """
+
+    def __init__(self, group_map: dict[str, str | None]) -> None:
+        self.group_map = group_map
+        self.type_to_ir: dict[TypeInfo, ClassIR] = {}
+        self.func_to_decl: dict[SymbolNode, FuncDecl] = {}
+
+    def type_to_rtype(self, typ: Type | None) -> RType:
+        if typ is None:
+            return object_rprimitive
+
+        typ = get_proper_type(typ)
+        if isinstance(typ, Instance):
+            if typ.type.fullname == "builtins.int":
+                return int_rprimitive
+            elif typ.type.fullname == "builtins.float":
+                return float_rprimitive
+            elif typ.type.fullname == "builtins.bool":
+                return bool_rprimitive
+            elif typ.type.fullname == "builtins.str":
+                return str_rprimitive
+            elif typ.type.fullname == "builtins.bytes":
+                return bytes_rprimitive
+            elif typ.type.fullname == "builtins.list":
+                return list_rprimitive
+            # Dict subclasses are at least somewhat common and we
+            # specifically support them, so make sure that dict operations
+            # get optimized on them.
+            elif any(cls.fullname == "builtins.dict" for cls in typ.type.mro):
+                return dict_rprimitive
+            elif typ.type.fullname == "builtins.set":
+                return set_rprimitive
+            elif typ.type.fullname == "builtins.tuple":
+                return tuple_rprimitive  # Varying-length tuple
+            elif typ.type.fullname == "builtins.range":
+                return range_rprimitive
+            elif typ.type in self.type_to_ir:
+                inst = RInstance(self.type_to_ir[typ.type])
+                # Treat protocols as Union[protocol, object], so that we can do fast
+                # method calls in the cases where the protocol is explicitly inherited from
+                # and fall back to generic operations when it isn't.
+                if typ.type.is_protocol:
+                    return RUnion([inst, object_rprimitive])
+                else:
+                    return inst
+            elif typ.type.fullname == "mypy_extensions.i64":
+                return int64_rprimitive
+            elif typ.type.fullname == "mypy_extensions.i32":
+                return int32_rprimitive
+            elif typ.type.fullname == "mypy_extensions.i16":
+                return int16_rprimitive
+            elif typ.type.fullname == "mypy_extensions.u8":
+                return uint8_rprimitive
+            else:
+                return object_rprimitive
+        elif isinstance(typ, TupleType):
+            # Use our unboxed tuples for raw tuples but fall back to
+            # being boxed for NamedTuple.
+            if typ.partial_fallback.type.fullname == "builtins.tuple":
+                return RTuple([self.type_to_rtype(t) for t in typ.items])
+            else:
+                return tuple_rprimitive
+        elif isinstance(typ, CallableType):
+            return object_rprimitive
+        elif isinstance(typ, NoneTyp):
+            return none_rprimitive
+        elif isinstance(typ, UnionType):
+            return RUnion.make_simplified_union([self.type_to_rtype(item) for item in typ.items])
+        elif isinstance(typ, AnyType):
+            return object_rprimitive
+        elif isinstance(typ, TypeType):
+            return object_rprimitive
+        elif isinstance(typ, TypeVarType):
+            # Erase type variable to upper bound.
+            # TODO: Erase to union if object has value restriction?
+            return self.type_to_rtype(typ.upper_bound)
+        elif isinstance(typ, PartialType):
+            assert typ.var.type is not None
+            return self.type_to_rtype(typ.var.type)
+        elif isinstance(typ, Overloaded):
+            return object_rprimitive
+        elif isinstance(typ, TypedDictType):
+            return dict_rprimitive
+        elif isinstance(typ, LiteralType):
+            return self.type_to_rtype(typ.fallback)
+        elif isinstance(typ, (UninhabitedType, UnboundType)):
+            # Sure, whatever!
+            return object_rprimitive
+
+        # I think we've covered everything that is supposed to
+        # actually show up, so anything else is a bug somewhere.
+        assert False, "unexpected type %s" % type(typ)
+
+    def get_arg_rtype(self, typ: Type, kind: ArgKind) -> RType:
+        if kind == ARG_STAR:
+            return tuple_rprimitive
+        elif kind == ARG_STAR2:
+            return dict_rprimitive
+        else:
+            return self.type_to_rtype(typ)
+
+    def fdef_to_sig(self, fdef: FuncDef) -> FuncSignature:
+        if isinstance(fdef.type, CallableType):
+            arg_types = [
+                self.get_arg_rtype(typ, kind)
+                for typ, kind in zip(fdef.type.arg_types, fdef.type.arg_kinds)
+            ]
+            arg_pos_onlys = [name is None for name in fdef.type.arg_names]
+            ret = self.type_to_rtype(fdef.type.ret_type)
+        else:
+            # Handle unannotated functions
+            arg_types = [object_rprimitive for _ in fdef.arguments]
+            arg_pos_onlys = [arg.pos_only for arg in fdef.arguments]
+            # We at least know the return type for __init__ methods will be None.
+            is_init_method = fdef.name == "__init__" and bool(fdef.info)
+            if is_init_method:
+                ret = none_rprimitive
+            else:
+                ret = object_rprimitive
+
+        # mypyc FuncSignatures (unlike mypy types) want to have a name
+        # present even when the argument is position only, since it is
+        # the sole way that FuncDecl arguments are tracked. This is
+        # generally fine except in some cases (like for computing
+        # init_sig) we need to produce FuncSignatures from a
+        # deserialized FuncDef that lacks arguments. We won't ever
+        # need to use those inside of a FuncIR, so we just make up
+        # some crap.
+        if hasattr(fdef, "arguments"):
+            arg_names = [arg.variable.name for arg in fdef.arguments]
+        else:
+            arg_names = [name or "" for name in fdef.arg_names]
+
+        args = [
+            RuntimeArg(arg_name, arg_type, arg_kind, arg_pos_only)
+            for arg_name, arg_kind, arg_type, arg_pos_only in zip(
+                arg_names, fdef.arg_kinds, arg_types, arg_pos_onlys
+            )
+        ]
+
+        # We force certain dunder methods to return objects to support letting them
+        # return NotImplemented. It also avoids some pointless boxing and unboxing,
+        # since tp_richcompare needs an object anyways.
+        if fdef.name in ("__eq__", "__ne__", "__lt__", "__gt__", "__le__", "__ge__"):
+            ret = object_rprimitive
+        return FuncSignature(args, ret)
+
+    def is_native_module(self, module: str) -> bool:
+        """Is the given module one compiled by mypyc?"""
+        return module in self.group_map
+
+    def is_native_ref_expr(self, expr: RefExpr) -> bool:
+        if expr.node is None:
+            return False
+        if "." in expr.node.fullname:
+            return self.is_native_module(expr.node.fullname.rpartition(".")[0])
+        return True
+
+    def is_native_module_ref_expr(self, expr: RefExpr) -> bool:
+        return self.is_native_ref_expr(expr) and expr.kind == GDEF

venv/lib/python3.12/site-packages/mypyc/irbuild/match.py

@@ -0,0 +1,355 @@
+from contextlib import contextmanager
+from typing import Generator, List, Optional, Tuple
+
+from mypy.nodes import MatchStmt, NameExpr, TypeInfo
+from mypy.patterns import (
+    AsPattern,
+    ClassPattern,
+    MappingPattern,
+    OrPattern,
+    Pattern,
+    SequencePattern,
+    SingletonPattern,
+    StarredPattern,
+    ValuePattern,
+)
+from mypy.traverser import TraverserVisitor
+from mypy.types import Instance, TupleType, get_proper_type
+from mypyc.ir.ops import BasicBlock, Value
+from mypyc.ir.rtypes import object_rprimitive
+from mypyc.irbuild.builder import IRBuilder
+from mypyc.primitives.dict_ops import (
+    dict_copy,
+    dict_del_item,
+    mapping_has_key,
+    supports_mapping_protocol,
+)
+from mypyc.primitives.generic_ops import generic_ssize_t_len_op
+from mypyc.primitives.list_ops import (
+    sequence_get_item,
+    sequence_get_slice,
+    supports_sequence_protocol,
+)
+from mypyc.primitives.misc_ops import fast_isinstance_op, slow_isinstance_op
+
+# From: https://peps.python.org/pep-0634/#class-patterns
+MATCHABLE_BUILTINS = {
+    "builtins.bool",
+    "builtins.bytearray",
+    "builtins.bytes",
+    "builtins.dict",
+    "builtins.float",
+    "builtins.frozenset",
+    "builtins.int",
+    "builtins.list",
+    "builtins.set",
+    "builtins.str",
+    "builtins.tuple",
+}
+
+
+class MatchVisitor(TraverserVisitor):
+    builder: IRBuilder
+    code_block: BasicBlock
+    next_block: BasicBlock
+    final_block: BasicBlock
+    subject: Value
+    match: MatchStmt
+
+    as_pattern: Optional[AsPattern] = None
+
+    def __init__(self, builder: IRBuilder, match_node: MatchStmt) -> None:
+        self.builder = builder
+
+        self.code_block = BasicBlock()
+        self.next_block = BasicBlock()
+        self.final_block = BasicBlock()
+
+        self.match = match_node
+        self.subject = builder.accept(match_node.subject)
+
+    def build_match_body(self, index: int) -> None:
+        self.builder.activate_block(self.code_block)
+
+        guard = self.match.guards[index]
+
+        if guard:
+            self.code_block = BasicBlock()
+
+            cond = self.builder.accept(guard)
+            self.builder.add_bool_branch(cond, self.code_block, self.next_block)
+
+            self.builder.activate_block(self.code_block)
+
+        self.builder.accept(self.match.bodies[index])
+        self.builder.goto(self.final_block)
+
+    def visit_match_stmt(self, m: MatchStmt) -> None:
+        for i, pattern in enumerate(m.patterns):
+            self.code_block = BasicBlock()
+            self.next_block = BasicBlock()
+
+            pattern.accept(self)
+
+            self.build_match_body(i)
+            self.builder.activate_block(self.next_block)
+
+        self.builder.goto_and_activate(self.final_block)
+
+    def visit_value_pattern(self, pattern: ValuePattern) -> None:
+        value = self.builder.accept(pattern.expr)
+
+        cond = self.builder.binary_op(self.subject, value, "==", pattern.expr.line)
+
+        self.bind_as_pattern(value)
+
+        self.builder.add_bool_branch(cond, self.code_block, self.next_block)
+
+    def visit_or_pattern(self, pattern: OrPattern) -> None:
+        backup_block = self.next_block
+        self.next_block = BasicBlock()
+
+        for p in pattern.patterns:
+            # Hack to ensure the as pattern is bound to each pattern in the
+            # "or" pattern, but not every subpattern
+            backup = self.as_pattern
+            p.accept(self)
+            self.as_pattern = backup
+
+            self.builder.activate_block(self.next_block)
+            self.next_block = BasicBlock()
+
+        self.next_block = backup_block
+        self.builder.goto(self.next_block)
+
+    def visit_class_pattern(self, pattern: ClassPattern) -> None:
+        # TODO: use faster instance check for native classes (while still
+        # making sure to account for inheritence)
+        isinstance_op = (
+            fast_isinstance_op
+            if self.builder.is_builtin_ref_expr(pattern.class_ref)
+            else slow_isinstance_op
+        )
+
+        cond = self.builder.call_c(
+            isinstance_op, [self.subject, self.builder.accept(pattern.class_ref)], pattern.line
+        )
+
+        self.builder.add_bool_branch(cond, self.code_block, self.next_block)
+
+        self.bind_as_pattern(self.subject, new_block=True)
+
+        if pattern.positionals:
+            if pattern.class_ref.fullname in MATCHABLE_BUILTINS:
+                self.builder.activate_block(self.code_block)
+                self.code_block = BasicBlock()
+
+                pattern.positionals[0].accept(self)
+
+                return
+
+            node = pattern.class_ref.node
+            assert isinstance(node, TypeInfo)
+
+            ty = node.names.get("__match_args__")
+            assert ty
+
+            match_args_type = get_proper_type(ty.type)
+            assert isinstance(match_args_type, TupleType)
+
+            match_args: List[str] = []
+
+            for item in match_args_type.items:
+                proper_item = get_proper_type(item)
+                assert isinstance(proper_item, Instance) and proper_item.last_known_value
+
+                match_arg = proper_item.last_known_value.value
+                assert isinstance(match_arg, str)
+
+                match_args.append(match_arg)
+
+            for i, expr in enumerate(pattern.positionals):
+                self.builder.activate_block(self.code_block)
+                self.code_block = BasicBlock()
+
+                # TODO: use faster "get_attr" method instead when calling on native or
+                # builtin objects
+                positional = self.builder.py_get_attr(self.subject, match_args[i], expr.line)
+
+                with self.enter_subpattern(positional):
+                    expr.accept(self)
+
+        for key, value in zip(pattern.keyword_keys, pattern.keyword_values):
+            self.builder.activate_block(self.code_block)
+            self.code_block = BasicBlock()
+
+            # TODO: same as above "get_attr" comment
+            attr = self.builder.py_get_attr(self.subject, key, value.line)
+
+            with self.enter_subpattern(attr):
+                value.accept(self)
+
+    def visit_as_pattern(self, pattern: AsPattern) -> None:
+        if pattern.pattern:
+            old_pattern = self.as_pattern
+            self.as_pattern = pattern
+            pattern.pattern.accept(self)
+            self.as_pattern = old_pattern
+
+        elif pattern.name:
+            target = self.builder.get_assignment_target(pattern.name)
+
+            self.builder.assign(target, self.subject, pattern.line)
+
+        self.builder.goto(self.code_block)
+
+    def visit_singleton_pattern(self, pattern: SingletonPattern) -> None:
+        if pattern.value is None:
+            obj = self.builder.none_object()
+        elif pattern.value is True:
+            obj = self.builder.true()
+        else:
+            obj = self.builder.false()
+
+        cond = self.builder.binary_op(self.subject, obj, "is", pattern.line)
+
+        self.builder.add_bool_branch(cond, self.code_block, self.next_block)
+
+    def visit_mapping_pattern(self, pattern: MappingPattern) -> None:
+        is_dict = self.builder.call_c(supports_mapping_protocol, [self.subject], pattern.line)
+
+        self.builder.add_bool_branch(is_dict, self.code_block, self.next_block)
+
+        keys: List[Value] = []
+
+        for key, value in zip(pattern.keys, pattern.values):
+            self.builder.activate_block(self.code_block)
+            self.code_block = BasicBlock()
+
+            key_value = self.builder.accept(key)
+            keys.append(key_value)
+
+            exists = self.builder.call_c(mapping_has_key, [self.subject, key_value], pattern.line)
+
+            self.builder.add_bool_branch(exists, self.code_block, self.next_block)
+            self.builder.activate_block(self.code_block)
+            self.code_block = BasicBlock()
+
+            item = self.builder.gen_method_call(
+                self.subject, "__getitem__", [key_value], object_rprimitive, pattern.line
+            )
+
+            with self.enter_subpattern(item):
+                value.accept(self)
+
+        if pattern.rest:
+            self.builder.activate_block(self.code_block)
+            self.code_block = BasicBlock()
+
+            rest = self.builder.call_c(dict_copy, [self.subject], pattern.rest.line)
+
+            target = self.builder.get_assignment_target(pattern.rest)
+
+            self.builder.assign(target, rest, pattern.rest.line)
+
+            for i, key_name in enumerate(keys):
+                self.builder.call_c(dict_del_item, [rest, key_name], pattern.keys[i].line)
+
+            self.builder.goto(self.code_block)
+
+    def visit_sequence_pattern(self, seq_pattern: SequencePattern) -> None:
+        star_index, capture, patterns = prep_sequence_pattern(seq_pattern)
+
+        is_list = self.builder.call_c(supports_sequence_protocol, [self.subject], seq_pattern.line)
+
+        self.builder.add_bool_branch(is_list, self.code_block, self.next_block)
+
+        self.builder.activate_block(self.code_block)
+        self.code_block = BasicBlock()
+
+        actual_len = self.builder.call_c(generic_ssize_t_len_op, [self.subject], seq_pattern.line)
+        min_len = len(patterns)
+
+        is_long_enough = self.builder.binary_op(
+            actual_len,
+            self.builder.load_int(min_len),
+            "==" if star_index is None else ">=",
+            seq_pattern.line,
+        )
+
+        self.builder.add_bool_branch(is_long_enough, self.code_block, self.next_block)
+
+        for i, pattern in enumerate(patterns):
+            self.builder.activate_block(self.code_block)
+            self.code_block = BasicBlock()
+
+            if star_index is not None and i >= star_index:
+                current = self.builder.binary_op(
+                    actual_len, self.builder.load_int(min_len - i), "-", pattern.line
+                )
+
+            else:
+                current = self.builder.load_int(i)
+
+            item = self.builder.call_c(sequence_get_item, [self.subject, current], pattern.line)
+
+            with self.enter_subpattern(item):
+                pattern.accept(self)
+
+        if capture and star_index is not None:
+            self.builder.activate_block(self.code_block)
+            self.code_block = BasicBlock()
+
+            capture_end = self.builder.binary_op(
+                actual_len, self.builder.load_int(min_len - star_index), "-", capture.line
+            )
+
+            rest = self.builder.call_c(
+                sequence_get_slice,
+                [self.subject, self.builder.load_int(star_index), capture_end],
+                capture.line,
+            )
+
+            target = self.builder.get_assignment_target(capture)
+            self.builder.assign(target, rest, capture.line)
+
+            self.builder.goto(self.code_block)
+
+    def bind_as_pattern(self, value: Value, new_block: bool = False) -> None:
+        if self.as_pattern and self.as_pattern.pattern and self.as_pattern.name:
+            if new_block:
+                self.builder.activate_block(self.code_block)
+                self.code_block = BasicBlock()
+
+            target = self.builder.get_assignment_target(self.as_pattern.name)
+            self.builder.assign(target, value, self.as_pattern.pattern.line)
+
+            self.as_pattern = None
+
+            if new_block:
+                self.builder.goto(self.code_block)
+
+    @contextmanager
+    def enter_subpattern(self, subject: Value) -> Generator[None, None, None]:
+        old_subject = self.subject
+        self.subject = subject
+        yield
+        self.subject = old_subject
+
+
+def prep_sequence_pattern(
+    seq_pattern: SequencePattern,
+) -> Tuple[Optional[int], Optional[NameExpr], List[Pattern]]:
+    star_index: Optional[int] = None
+    capture: Optional[NameExpr] = None
+    patterns: List[Pattern] = []
+
+    for i, pattern in enumerate(seq_pattern.patterns):
+        if isinstance(pattern, StarredPattern):
+            star_index = i
+            capture = pattern.capture
+
+        else:
+            patterns.append(pattern)
+
+    return star_index, capture, patterns

venv/lib/python3.12/site-packages/mypyc/irbuild/nonlocalcontrol.py

@@ -0,0 +1,198 @@
+"""Helpers for dealing with nonlocal control such as 'break' and 'return'.
+
+Model how these behave differently in different contexts.
+"""
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING
+
+from mypyc.ir.ops import (
+    NO_TRACEBACK_LINE_NO,
+    BasicBlock,
+    Branch,
+    Goto,
+    Integer,
+    Register,
+    Return,
+    Unreachable,
+    Value,
+)
+from mypyc.irbuild.targets import AssignmentTarget
+from mypyc.primitives.exc_ops import restore_exc_info_op, set_stop_iteration_value
+
+if TYPE_CHECKING:
+    from mypyc.irbuild.builder import IRBuilder
+
+
+class NonlocalControl:
+    """ABC representing a stack frame of constructs that modify nonlocal control flow.
+
+    The nonlocal control flow constructs are break, continue, and
+    return, and their behavior is modified by a number of other
+    constructs.  The most obvious is loop, which override where break
+    and continue jump to, but also `except` (which needs to clear
+    exc_info when left) and (eventually) finally blocks (which need to
+    ensure that the finally block is always executed when leaving the
+    try/except blocks).
+    """
+
+    @abstractmethod
+    def gen_break(self, builder: IRBuilder, line: int) -> None:
+        pass
+
+    @abstractmethod
+    def gen_continue(self, builder: IRBuilder, line: int) -> None:
+        pass
+
+    @abstractmethod
+    def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
+        pass
+
+
+class BaseNonlocalControl(NonlocalControl):
+    """Default nonlocal control outside any statements that affect it."""
+
+    def gen_break(self, builder: IRBuilder, line: int) -> None:
+        assert False, "break outside of loop"
+
+    def gen_continue(self, builder: IRBuilder, line: int) -> None:
+        assert False, "continue outside of loop"
+
+    def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
+        builder.add(Return(value))
+
+
+class LoopNonlocalControl(NonlocalControl):
+    """Nonlocal control within a loop."""
+
+    def __init__(
+        self, outer: NonlocalControl, continue_block: BasicBlock, break_block: BasicBlock
+    ) -> None:
+        self.outer = outer
+        self.continue_block = continue_block
+        self.break_block = break_block
+
+    def gen_break(self, builder: IRBuilder, line: int) -> None:
+        builder.add(Goto(self.break_block))
+
+    def gen_continue(self, builder: IRBuilder, line: int) -> None:
+        builder.add(Goto(self.continue_block))
+
+    def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
+        self.outer.gen_return(builder, value, line)
+
+
+class GeneratorNonlocalControl(BaseNonlocalControl):
+    """Default nonlocal control in a generator function outside statements."""
+
+    def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
+        # Assign an invalid next label number so that the next time
+        # __next__ is called, we jump to the case in which
+        # StopIteration is raised.
+        builder.assign(builder.fn_info.generator_class.next_label_target, Integer(-1), line)
+
+        # Raise a StopIteration containing a field for the value that
+        # should be returned. Before doing so, create a new block
+        # without an error handler set so that the implicitly thrown
+        # StopIteration isn't caught by except blocks inside of the
+        # generator function.
+        builder.builder.push_error_handler(None)
+        builder.goto_and_activate(BasicBlock())
+
+        # Skip creating a traceback frame when we raise here, because
+        # we don't care about the traceback frame and it is kind of
+        # expensive since raising StopIteration is an extremely common
+        # case.  Also we call a special internal function to set
+        # StopIteration instead of using RaiseStandardError because
+        # the obvious thing doesn't work if the value is a tuple
+        # (???).
+        builder.call_c(set_stop_iteration_value, [value], NO_TRACEBACK_LINE_NO)
+        builder.add(Unreachable())
+        builder.builder.pop_error_handler()
+
+
+class CleanupNonlocalControl(NonlocalControl):
+    """Abstract nonlocal control that runs some cleanup code."""
+
+    def __init__(self, outer: NonlocalControl) -> None:
+        self.outer = outer
+
+    @abstractmethod
+    def gen_cleanup(self, builder: IRBuilder, line: int) -> None:
+        ...
+
+    def gen_break(self, builder: IRBuilder, line: int) -> None:
+        self.gen_cleanup(builder, line)
+        self.outer.gen_break(builder, line)
+
+    def gen_continue(self, builder: IRBuilder, line: int) -> None:
+        self.gen_cleanup(builder, line)
+        self.outer.gen_continue(builder, line)
+
+    def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
+        self.gen_cleanup(builder, line)
+        self.outer.gen_return(builder, value, line)
+
+
+class TryFinallyNonlocalControl(NonlocalControl):
+    """Nonlocal control within try/finally."""
+
+    def __init__(self, target: BasicBlock) -> None:
+        self.target = target
+        self.ret_reg: None | Register | AssignmentTarget = None
+
+    def gen_break(self, builder: IRBuilder, line: int) -> None:
+        builder.error("break inside try/finally block is unimplemented", line)
+
+    def gen_continue(self, builder: IRBuilder, line: int) -> None:
+        builder.error("continue inside try/finally block is unimplemented", line)
+
+    def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
+        if self.ret_reg is None:
+            if builder.fn_info.is_generator:
+                self.ret_reg = builder.make_spill_target(builder.ret_types[-1])
+            else:
+                self.ret_reg = Register(builder.ret_types[-1])
+        # assert needed because of apparent mypy bug... it loses track of the union
+        # and infers the type as object
+        assert isinstance(self.ret_reg, (Register, AssignmentTarget))
+        builder.assign(self.ret_reg, value, line)
+
+        builder.add(Goto(self.target))
+
+
+class ExceptNonlocalControl(CleanupNonlocalControl):
+    """Nonlocal control for except blocks.
+
+    Just makes sure that sys.exc_info always gets restored when we leave.
+    This is super annoying.
+    """
+
+    def __init__(self, outer: NonlocalControl, saved: Value | AssignmentTarget) -> None:
+        super().__init__(outer)
+        self.saved = saved
+
+    def gen_cleanup(self, builder: IRBuilder, line: int) -> None:
+        builder.call_c(restore_exc_info_op, [builder.read(self.saved)], line)
+
+
+class FinallyNonlocalControl(CleanupNonlocalControl):
+    """Nonlocal control for finally blocks.
+
+    Just makes sure that sys.exc_info always gets restored when we
+    leave and the return register is decrefed if it isn't null.
+    """
+
+    def __init__(self, outer: NonlocalControl, saved: Value) -> None:
+        super().__init__(outer)
+        self.saved = saved
+
+    def gen_cleanup(self, builder: IRBuilder, line: int) -> None:
+        # Restore the old exc_info
+        target, cleanup = BasicBlock(), BasicBlock()
+        builder.add(Branch(self.saved, target, cleanup, Branch.IS_ERROR))
+        builder.activate_block(cleanup)
+        builder.call_c(restore_exc_info_op, [self.saved], line)
+        builder.goto_and_activate(target)

venv/lib/python3.12/site-packages/mypyc/irbuild/prebuildvisitor.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+from mypy.nodes import (
+    Block,
+    Decorator,
+    Expression,
+    FuncDef,
+    FuncItem,
+    Import,
+    LambdaExpr,
+    MemberExpr,
+    MypyFile,
+    NameExpr,
+    Node,
+    SymbolNode,
+    Var,
+)
+from mypy.traverser import ExtendedTraverserVisitor
+from mypyc.errors import Errors
+
+
+class PreBuildVisitor(ExtendedTraverserVisitor):
+    """Mypy file AST visitor run before building the IR.
+
+    This collects various things, including:
+
+    * Determine relationships between nested functions and functions that
+      contain nested functions
+    * Find non-local variables (free variables)
+    * Find property setters
+    * Find decorators of functions
+    * Find module import groups
+
+    The main IR build pass uses this information.
+    """
+
+    def __init__(
+        self,
+        errors: Errors,
+        current_file: MypyFile,
+        decorators_to_remove: dict[FuncDef, list[int]],
+    ) -> None:
+        super().__init__()
+        # Dict from a function to symbols defined directly in the
+        # function that are used as non-local (free) variables within a
+        # nested function.
+        self.free_variables: dict[FuncItem, set[SymbolNode]] = {}
+
+        # Intermediate data structure used to find the function where
+        # a SymbolNode is declared. Initially this may point to a
+        # function nested inside the function with the declaration,
+        # but we'll eventually update this to refer to the function
+        # with the declaration.
+        self.symbols_to_funcs: dict[SymbolNode, FuncItem] = {}
+
+        # Stack representing current function nesting.
+        self.funcs: list[FuncItem] = []
+
+        # All property setters encountered so far.
+        self.prop_setters: set[FuncDef] = set()
+
+        # A map from any function that contains nested functions to
+        # a set of all the functions that are nested within it.
+        self.encapsulating_funcs: dict[FuncItem, list[FuncItem]] = {}
+
+        # Map nested function to its parent/encapsulating function.
+        self.nested_funcs: dict[FuncItem, FuncItem] = {}
+
+        # Map function to its non-special decorators.
+        self.funcs_to_decorators: dict[FuncDef, list[Expression]] = {}
+
+        # Map function to indices of decorators to remove
+        self.decorators_to_remove: dict[FuncDef, list[int]] = decorators_to_remove
+
+        # A mapping of import groups (a series of Import nodes with
+        # nothing inbetween) where each group is keyed by its first
+        # import node.
+        self.module_import_groups: dict[Import, list[Import]] = {}
+        self._current_import_group: Import | None = None
+
+        self.errors: Errors = errors
+
+        self.current_file: MypyFile = current_file
+
+    def visit(self, o: Node) -> bool:
+        if not isinstance(o, Import):
+            self._current_import_group = None
+        return True
+
+    def visit_block(self, block: Block) -> None:
+        self._current_import_group = None
+        super().visit_block(block)
+        self._current_import_group = None
+
+    def visit_decorator(self, dec: Decorator) -> None:
+        if dec.decorators:
+            # Only add the function being decorated if there exist
+            # (ordinary) decorators in the decorator list. Certain
+            # decorators (such as @property, @abstractmethod) are
+            # special cased and removed from this list by
+            # mypy. Functions decorated only by special decorators
+            # (and property setters) are not treated as decorated
+            # functions by the IR builder.
+            if isinstance(dec.decorators[0], MemberExpr) and dec.decorators[0].name == "setter":
+                # Property setters are not treated as decorated methods.
+                self.prop_setters.add(dec.func)
+            else:
+                decorators_to_store = dec.decorators.copy()
+                if dec.func in self.decorators_to_remove:
+                    to_remove = self.decorators_to_remove[dec.func]
+
+                    for i in reversed(to_remove):
+                        del decorators_to_store[i]
+                    # if all of the decorators are removed, we shouldn't treat this as a decorated
+                    # function because there aren't any decorators to apply
+                    if not decorators_to_store:
+                        return
+
+                self.funcs_to_decorators[dec.func] = decorators_to_store
+        super().visit_decorator(dec)
+
+    def visit_func_def(self, fdef: FuncItem) -> None:
+        # TODO: What about overloaded functions?
+        self.visit_func(fdef)
+
+    def visit_lambda_expr(self, expr: LambdaExpr) -> None:
+        self.visit_func(expr)
+
+    def visit_func(self, func: FuncItem) -> None:
+        # If there were already functions or lambda expressions
+        # defined in the function stack, then note the previous
+        # FuncItem as containing a nested function and the current
+        # FuncItem as being a nested function.
+        if self.funcs:
+            # Add the new func to the set of nested funcs within the
+            # func at top of the func stack.
+            self.encapsulating_funcs.setdefault(self.funcs[-1], []).append(func)
+            # Add the func at top of the func stack as the parent of
+            # new func.
+            self.nested_funcs[func] = self.funcs[-1]
+
+        self.funcs.append(func)
+        super().visit_func(func)
+        self.funcs.pop()
+
+    def visit_import(self, imp: Import) -> None:
+        if self._current_import_group is not None:
+            self.module_import_groups[self._current_import_group].append(imp)
+        else:
+            self.module_import_groups[imp] = [imp]
+            self._current_import_group = imp
+        super().visit_import(imp)
+
+    def visit_name_expr(self, expr: NameExpr) -> None:
+        if isinstance(expr.node, (Var, FuncDef)):
+            self.visit_symbol_node(expr.node)
+
+    def visit_var(self, var: Var) -> None:
+        self.visit_symbol_node(var)
+
+    def visit_symbol_node(self, symbol: SymbolNode) -> None:
+        if not self.funcs:
+            # We are not inside a function and hence do not need to do
+            # anything regarding free variables.
+            return
+
+        if symbol in self.symbols_to_funcs:
+            orig_func = self.symbols_to_funcs[symbol]
+            if self.is_parent(self.funcs[-1], orig_func):
+                # The function in which the symbol was previously seen is
+                # nested within the function currently being visited. Thus
+                # the current function is a better candidate to contain the
+                # declaration.
+                self.symbols_to_funcs[symbol] = self.funcs[-1]
+                # TODO: Remove from the orig_func free_variables set?
+                self.free_variables.setdefault(self.funcs[-1], set()).add(symbol)
+
+            elif self.is_parent(orig_func, self.funcs[-1]):
+                # The SymbolNode instance has already been visited
+                # before in a parent function, thus it's a non-local
+                # symbol.
+                self.add_free_variable(symbol)
+
+        else:
+            # This is the first time the SymbolNode is being
+            # visited. We map the SymbolNode to the current FuncDef
+            # being visited to note where it was first visited.
+            self.symbols_to_funcs[symbol] = self.funcs[-1]
+
+    def is_parent(self, fitem: FuncItem, child: FuncItem) -> bool:
+        # Check if child is nested within fdef (possibly indirectly
+        # within multiple nested functions).
+        if child not in self.nested_funcs:
+            return False
+        parent = self.nested_funcs[child]
+        return parent == fitem or self.is_parent(fitem, parent)
+
+    def add_free_variable(self, symbol: SymbolNode) -> None:
+        # Find the function where the symbol was (likely) first declared,
+        # and mark is as a non-local symbol within that function.
+        func = self.symbols_to_funcs[symbol]
+        self.free_variables.setdefault(func, set()).add(symbol)

venv/lib/python3.12/site-packages/mypyc/irbuild/prepare.py

@@ -0,0 +1,609 @@
+"""Prepare for IR transform.
+
+This needs to run after type checking and before generating IR.
+
+For example, construct partially initialized FuncIR and ClassIR
+objects for all functions and classes. This allows us to bind
+references to functions and classes before we've generated full IR for
+functions or classes.  The actual IR transform will then populate all
+the missing bits, such as function bodies (basic blocks).
+
+Also build a mapping from mypy TypeInfos to ClassIR objects.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Iterable, NamedTuple, Tuple
+
+from mypy.build import Graph
+from mypy.nodes import (
+    ARG_STAR,
+    ARG_STAR2,
+    CallExpr,
+    ClassDef,
+    Decorator,
+    Expression,
+    FuncDef,
+    MemberExpr,
+    MypyFile,
+    NameExpr,
+    OverloadedFuncDef,
+    RefExpr,
+    SymbolNode,
+    TypeInfo,
+    Var,
+)
+from mypy.semanal import refers_to_fullname
+from mypy.traverser import TraverserVisitor
+from mypy.types import Instance, Type, get_proper_type
+from mypyc.common import PROPSET_PREFIX, get_id_from_name
+from mypyc.crash import catch_errors
+from mypyc.errors import Errors
+from mypyc.ir.class_ir import ClassIR
+from mypyc.ir.func_ir import (
+    FUNC_CLASSMETHOD,
+    FUNC_NORMAL,
+    FUNC_STATICMETHOD,
+    FuncDecl,
+    FuncSignature,
+    RuntimeArg,
+)
+from mypyc.ir.ops import DeserMaps
+from mypyc.ir.rtypes import RInstance, RType, dict_rprimitive, none_rprimitive, tuple_rprimitive
+from mypyc.irbuild.mapper import Mapper
+from mypyc.irbuild.util import (
+    get_func_def,
+    get_mypyc_attrs,
+    is_dataclass,
+    is_extension_class,
+    is_trait,
+)
+from mypyc.options import CompilerOptions
+from mypyc.sametype import is_same_type
+
+
+def build_type_map(
+    mapper: Mapper,
+    modules: list[MypyFile],
+    graph: Graph,
+    types: dict[Expression, Type],
+    options: CompilerOptions,
+    errors: Errors,
+) -> None:
+    # Collect all classes defined in everything we are compiling
+    classes = []
+    for module in modules:
+        module_classes = [node for node in module.defs if isinstance(node, ClassDef)]
+        classes.extend([(module, cdef) for cdef in module_classes])
+
+    # Collect all class mappings so that we can bind arbitrary class name
+    # references even if there are import cycles.
+    for module, cdef in classes:
+        class_ir = ClassIR(
+            cdef.name, module.fullname, is_trait(cdef), is_abstract=cdef.info.is_abstract
+        )
+        class_ir.is_ext_class = is_extension_class(cdef)
+        if class_ir.is_ext_class:
+            class_ir.deletable = cdef.info.deletable_attributes.copy()
+        # If global optimizations are disabled, turn of tracking of class children
+        if not options.global_opts:
+            class_ir.children = None
+        mapper.type_to_ir[cdef.info] = class_ir
+
+    # Populate structural information in class IR for extension classes.
+    for module, cdef in classes:
+        with catch_errors(module.path, cdef.line):
+            if mapper.type_to_ir[cdef.info].is_ext_class:
+                prepare_class_def(module.path, module.fullname, cdef, errors, mapper)
+            else:
+                prepare_non_ext_class_def(module.path, module.fullname, cdef, errors, mapper)
+
+    # Prepare implicit attribute accessors as needed if an attribute overrides a property.
+    for module, cdef in classes:
+        class_ir = mapper.type_to_ir[cdef.info]
+        if class_ir.is_ext_class:
+            prepare_implicit_property_accessors(cdef.info, class_ir, module.fullname, mapper)
+
+    # Collect all the functions also. We collect from the symbol table
+    # so that we can easily pick out the right copy of a function that
+    # is conditionally defined.
+    for module in modules:
+        for func in get_module_func_defs(module):
+            prepare_func_def(module.fullname, None, func, mapper)
+            # TODO: what else?
+
+    # Check for incompatible attribute definitions that were not
+    # flagged by mypy but can't be supported when compiling.
+    for module, cdef in classes:
+        class_ir = mapper.type_to_ir[cdef.info]
+        for attr in class_ir.attributes:
+            for base_ir in class_ir.mro[1:]:
+                if attr in base_ir.attributes:
+                    if not is_same_type(class_ir.attributes[attr], base_ir.attributes[attr]):
+                        node = cdef.info.names[attr].node
+                        assert node is not None
+                        kind = "trait" if base_ir.is_trait else "class"
+                        errors.error(
+                            f'Type of "{attr}" is incompatible with '
+                            f'definition in {kind} "{base_ir.name}"',
+                            module.path,
+                            node.line,
+                        )
+
+
+def is_from_module(node: SymbolNode, module: MypyFile) -> bool:
+    return node.fullname == module.fullname + "." + node.name
+
+
+def load_type_map(mapper: Mapper, modules: list[MypyFile], deser_ctx: DeserMaps) -> None:
+    """Populate a Mapper with deserialized IR from a list of modules."""
+    for module in modules:
+        for name, node in module.names.items():
+            if isinstance(node.node, TypeInfo) and is_from_module(node.node, module):
+                ir = deser_ctx.classes[node.node.fullname]
+                mapper.type_to_ir[node.node] = ir
+                mapper.func_to_decl[node.node] = ir.ctor
+
+    for module in modules:
+        for func in get_module_func_defs(module):
+            func_id = get_id_from_name(func.name, func.fullname, func.line)
+            mapper.func_to_decl[func] = deser_ctx.functions[func_id].decl
+
+
+def get_module_func_defs(module: MypyFile) -> Iterable[FuncDef]:
+    """Collect all of the (non-method) functions declared in a module."""
+    for name, node in module.names.items():
+        # We need to filter out functions that are imported or
+        # aliases.  The best way to do this seems to be by
+        # checking that the fullname matches.
+        if isinstance(node.node, (FuncDef, Decorator, OverloadedFuncDef)) and is_from_module(
+            node.node, module
+        ):
+            yield get_func_def(node.node)
+
+
+def prepare_func_def(
+    module_name: str, class_name: str | None, fdef: FuncDef, mapper: Mapper
+) -> FuncDecl:
+    kind = (
+        FUNC_STATICMETHOD
+        if fdef.is_static
+        else (FUNC_CLASSMETHOD if fdef.is_class else FUNC_NORMAL)
+    )
+    decl = FuncDecl(fdef.name, class_name, module_name, mapper.fdef_to_sig(fdef), kind)
+    mapper.func_to_decl[fdef] = decl
+    return decl
+
+
+def prepare_method_def(
+    ir: ClassIR, module_name: str, cdef: ClassDef, mapper: Mapper, node: FuncDef | Decorator
+) -> None:
+    if isinstance(node, FuncDef):
+        ir.method_decls[node.name] = prepare_func_def(module_name, cdef.name, node, mapper)
+    elif isinstance(node, Decorator):
+        # TODO: do something about abstract methods here. Currently, they are handled just like
+        # normal methods.
+        decl = prepare_func_def(module_name, cdef.name, node.func, mapper)
+        if not node.decorators:
+            ir.method_decls[node.name] = decl
+        elif isinstance(node.decorators[0], MemberExpr) and node.decorators[0].name == "setter":
+            # Make property setter name different than getter name so there are no
+            # name clashes when generating C code, and property lookup at the IR level
+            # works correctly.
+            decl.name = PROPSET_PREFIX + decl.name
+            decl.is_prop_setter = True
+            # Making the argument implicitly positional-only avoids unnecessary glue methods
+            decl.sig.args[1].pos_only = True
+            ir.method_decls[PROPSET_PREFIX + node.name] = decl
+
+        if node.func.is_property:
+            assert node.func.type, f"Expected return type annotation for property '{node.name}'"
+            decl.is_prop_getter = True
+            ir.property_types[node.name] = decl.sig.ret_type
+
+
+def is_valid_multipart_property_def(prop: OverloadedFuncDef) -> bool:
+    # Checks to ensure supported property decorator semantics
+    if len(prop.items) != 2:
+        return False
+
+    getter = prop.items[0]
+    setter = prop.items[1]
+
+    return (
+        isinstance(getter, Decorator)
+        and isinstance(setter, Decorator)
+        and getter.func.is_property
+        and len(setter.decorators) == 1
+        and isinstance(setter.decorators[0], MemberExpr)
+        and setter.decorators[0].name == "setter"
+    )
+
+
+def can_subclass_builtin(builtin_base: str) -> bool:
+    # BaseException and dict are special cased.
+    return builtin_base in (
+        (
+            "builtins.Exception",
+            "builtins.LookupError",
+            "builtins.IndexError",
+            "builtins.Warning",
+            "builtins.UserWarning",
+            "builtins.ValueError",
+            "builtins.object",
+        )
+    )
+
+
+def prepare_class_def(
+    path: str, module_name: str, cdef: ClassDef, errors: Errors, mapper: Mapper
+) -> None:
+    """Populate the interface-level information in a class IR.
+
+    This includes attribute and method declarations, and the MRO, among other things, but
+    method bodies are generated in a later pass.
+    """
+
+    ir = mapper.type_to_ir[cdef.info]
+    info = cdef.info
+
+    attrs = get_mypyc_attrs(cdef)
+    if attrs.get("allow_interpreted_subclasses") is True:
+        ir.allow_interpreted_subclasses = True
+    if attrs.get("serializable") is True:
+        # Supports copy.copy and pickle (including subclasses)
+        ir._serializable = True
+
+    # Check for subclassing from builtin types
+    for cls in info.mro:
+        # Special case exceptions and dicts
+        # XXX: How do we handle *other* things??
+        if cls.fullname == "builtins.BaseException":
+            ir.builtin_base = "PyBaseExceptionObject"
+        elif cls.fullname == "builtins.dict":
+            ir.builtin_base = "PyDictObject"
+        elif cls.fullname.startswith("builtins."):
+            if not can_subclass_builtin(cls.fullname):
+                # Note that if we try to subclass a C extension class that
+                # isn't in builtins, bad things will happen and we won't
+                # catch it here! But this should catch a lot of the most
+                # common pitfalls.
+                errors.error(
+                    "Inheriting from most builtin types is unimplemented", path, cdef.line
+                )
+
+    # Set up the parent class
+    bases = [mapper.type_to_ir[base.type] for base in info.bases if base.type in mapper.type_to_ir]
+    if len(bases) > 1 and any(not c.is_trait for c in bases) and bases[0].is_trait:
+        # If the first base is a non-trait, don't ever error here. While it is correct
+        # to error if a trait comes before the next non-trait base (e.g. non-trait, trait,
+        # non-trait), it's pointless, confusing noise from the bigger issue: multiple
+        # inheritance is *not* supported.
+        errors.error("Non-trait base must appear first in parent list", path, cdef.line)
+    ir.traits = [c for c in bases if c.is_trait]
+
+    mro = []  # All mypyc base classes
+    base_mro = []  # Non-trait mypyc base classes
+    for cls in info.mro:
+        if cls not in mapper.type_to_ir:
+            if cls.fullname != "builtins.object":
+                ir.inherits_python = True
+            continue
+        base_ir = mapper.type_to_ir[cls]
+        if not base_ir.is_trait:
+            base_mro.append(base_ir)
+        mro.append(base_ir)
+
+        if cls.defn.removed_base_type_exprs or not base_ir.is_ext_class:
+            ir.inherits_python = True
+
+    base_idx = 1 if not ir.is_trait else 0
+    if len(base_mro) > base_idx:
+        ir.base = base_mro[base_idx]
+    ir.mro = mro
+    ir.base_mro = base_mro
+
+    prepare_methods_and_attributes(cdef, ir, path, module_name, errors, mapper)
+    prepare_init_method(cdef, ir, module_name, mapper)
+
+    for base in bases:
+        if base.children is not None:
+            base.children.append(ir)
+
+    if is_dataclass(cdef):
+        ir.is_augmented = True
+
+
+def prepare_methods_and_attributes(
+    cdef: ClassDef, ir: ClassIR, path: str, module_name: str, errors: Errors, mapper: Mapper
+) -> None:
+    """Populate attribute and method declarations."""
+    info = cdef.info
+    for name, node in info.names.items():
+        # Currently all plugin generated methods are dummies and not included.
+        if node.plugin_generated:
+            continue
+
+        if isinstance(node.node, Var):
+            assert node.node.type, "Class member %s missing type" % name
+            if not node.node.is_classvar and name not in ("__slots__", "__deletable__"):
+                attr_rtype = mapper.type_to_rtype(node.node.type)
+                if ir.is_trait and attr_rtype.error_overlap:
+                    # Traits don't have attribute definedness bitmaps, so use
+                    # property accessor methods to access attributes that need them.
+                    # We will generate accessor implementations that use the class bitmap
+                    # for any concrete subclasses.
+                    add_getter_declaration(ir, name, attr_rtype, module_name)
+                    add_setter_declaration(ir, name, attr_rtype, module_name)
+                ir.attributes[name] = attr_rtype
+        elif isinstance(node.node, (FuncDef, Decorator)):
+            prepare_method_def(ir, module_name, cdef, mapper, node.node)
+        elif isinstance(node.node, OverloadedFuncDef):
+            # Handle case for property with both a getter and a setter
+            if node.node.is_property:
+                if is_valid_multipart_property_def(node.node):
+                    for item in node.node.items:
+                        prepare_method_def(ir, module_name, cdef, mapper, item)
+                else:
+                    errors.error("Unsupported property decorator semantics", path, cdef.line)
+
+            # Handle case for regular function overload
+            else:
+                assert node.node.impl
+                prepare_method_def(ir, module_name, cdef, mapper, node.node.impl)
+
+    if ir.builtin_base:
+        ir.attributes.clear()
+
+
+def prepare_implicit_property_accessors(
+    info: TypeInfo, ir: ClassIR, module_name: str, mapper: Mapper
+) -> None:
+    concrete_attributes = set()
+    for base in ir.base_mro:
+        for name, attr_rtype in base.attributes.items():
+            concrete_attributes.add(name)
+            add_property_methods_for_attribute_if_needed(
+                info, ir, name, attr_rtype, module_name, mapper
+            )
+    for base in ir.mro[1:]:
+        if base.is_trait:
+            for name, attr_rtype in base.attributes.items():
+                if name not in concrete_attributes:
+                    add_property_methods_for_attribute_if_needed(
+                        info, ir, name, attr_rtype, module_name, mapper
+                    )
+
+
+def add_property_methods_for_attribute_if_needed(
+    info: TypeInfo,
+    ir: ClassIR,
+    attr_name: str,
+    attr_rtype: RType,
+    module_name: str,
+    mapper: Mapper,
+) -> None:
+    """Add getter and/or setter for attribute if defined as property in a base class.
+
+    Only add declarations. The body IR will be synthesized later during irbuild.
+    """
+    for base in info.mro[1:]:
+        if base in mapper.type_to_ir:
+            base_ir = mapper.type_to_ir[base]
+            n = base.names.get(attr_name)
+            if n is None:
+                continue
+            node = n.node
+            if isinstance(node, Decorator) and node.name not in ir.method_decls:
+                # Defined as a read-only property in base class/trait
+                add_getter_declaration(ir, attr_name, attr_rtype, module_name)
+            elif isinstance(node, OverloadedFuncDef) and is_valid_multipart_property_def(node):
+                # Defined as a read-write property in base class/trait
+                add_getter_declaration(ir, attr_name, attr_rtype, module_name)
+                add_setter_declaration(ir, attr_name, attr_rtype, module_name)
+            elif base_ir.is_trait and attr_rtype.error_overlap:
+                add_getter_declaration(ir, attr_name, attr_rtype, module_name)
+                add_setter_declaration(ir, attr_name, attr_rtype, module_name)
+
+
+def add_getter_declaration(
+    ir: ClassIR, attr_name: str, attr_rtype: RType, module_name: str
+) -> None:
+    self_arg = RuntimeArg("self", RInstance(ir), pos_only=True)
+    sig = FuncSignature([self_arg], attr_rtype)
+    decl = FuncDecl(attr_name, ir.name, module_name, sig, FUNC_NORMAL)
+    decl.is_prop_getter = True
+    decl.implicit = True  # Triggers synthesization
+    ir.method_decls[attr_name] = decl
+    ir.property_types[attr_name] = attr_rtype  # TODO: Needed??
+
+
+def add_setter_declaration(
+    ir: ClassIR, attr_name: str, attr_rtype: RType, module_name: str
+) -> None:
+    self_arg = RuntimeArg("self", RInstance(ir), pos_only=True)
+    value_arg = RuntimeArg("value", attr_rtype, pos_only=True)
+    sig = FuncSignature([self_arg, value_arg], none_rprimitive)
+    setter_name = PROPSET_PREFIX + attr_name
+    decl = FuncDecl(setter_name, ir.name, module_name, sig, FUNC_NORMAL)
+    decl.is_prop_setter = True
+    decl.implicit = True  # Triggers synthesization
+    ir.method_decls[setter_name] = decl
+
+
+def prepare_init_method(cdef: ClassDef, ir: ClassIR, module_name: str, mapper: Mapper) -> None:
+    # Set up a constructor decl
+    init_node = cdef.info["__init__"].node
+    if not ir.is_trait and not ir.builtin_base and isinstance(init_node, FuncDef):
+        init_sig = mapper.fdef_to_sig(init_node)
+
+        defining_ir = mapper.type_to_ir.get(init_node.info)
+        # If there is a nontrivial __init__ that wasn't defined in an
+        # extension class, we need to make the constructor take *args,
+        # **kwargs so it can call tp_init.
+        if (
+            defining_ir is None
+            or not defining_ir.is_ext_class
+            or cdef.info["__init__"].plugin_generated
+        ) and init_node.info.fullname != "builtins.object":
+            init_sig = FuncSignature(
+                [
+                    init_sig.args[0],
+                    RuntimeArg("args", tuple_rprimitive, ARG_STAR),
+                    RuntimeArg("kwargs", dict_rprimitive, ARG_STAR2),
+                ],
+                init_sig.ret_type,
+            )
+
+        last_arg = len(init_sig.args) - init_sig.num_bitmap_args
+        ctor_sig = FuncSignature(init_sig.args[1:last_arg], RInstance(ir))
+        ir.ctor = FuncDecl(cdef.name, None, module_name, ctor_sig)
+        mapper.func_to_decl[cdef.info] = ir.ctor
+
+
+def prepare_non_ext_class_def(
+    path: str, module_name: str, cdef: ClassDef, errors: Errors, mapper: Mapper
+) -> None:
+    ir = mapper.type_to_ir[cdef.info]
+    info = cdef.info
+
+    for name, node in info.names.items():
+        if isinstance(node.node, (FuncDef, Decorator)):
+            prepare_method_def(ir, module_name, cdef, mapper, node.node)
+        elif isinstance(node.node, OverloadedFuncDef):
+            # Handle case for property with both a getter and a setter
+            if node.node.is_property:
+                if not is_valid_multipart_property_def(node.node):
+                    errors.error("Unsupported property decorator semantics", path, cdef.line)
+                for item in node.node.items:
+                    prepare_method_def(ir, module_name, cdef, mapper, item)
+            # Handle case for regular function overload
+            else:
+                prepare_method_def(ir, module_name, cdef, mapper, get_func_def(node.node))
+
+    if any(cls in mapper.type_to_ir and mapper.type_to_ir[cls].is_ext_class for cls in info.mro):
+        errors.error(
+            "Non-extension classes may not inherit from extension classes", path, cdef.line
+        )
+
+
+RegisterImplInfo = Tuple[TypeInfo, FuncDef]
+
+
+class SingledispatchInfo(NamedTuple):
+    singledispatch_impls: dict[FuncDef, list[RegisterImplInfo]]
+    decorators_to_remove: dict[FuncDef, list[int]]
+
+
+def find_singledispatch_register_impls(
+    modules: list[MypyFile], errors: Errors
+) -> SingledispatchInfo:
+    visitor = SingledispatchVisitor(errors)
+    for module in modules:
+        visitor.current_path = module.path
+        module.accept(visitor)
+    return SingledispatchInfo(visitor.singledispatch_impls, visitor.decorators_to_remove)
+
+
+class SingledispatchVisitor(TraverserVisitor):
+    current_path: str
+
+    def __init__(self, errors: Errors) -> None:
+        super().__init__()
+
+        # Map of main singledispatch function to list of registered implementations
+        self.singledispatch_impls: defaultdict[FuncDef, list[RegisterImplInfo]] = defaultdict(list)
+
+        # Map of decorated function to the indices of any decorators to remove
+        self.decorators_to_remove: dict[FuncDef, list[int]] = {}
+
+        self.errors: Errors = errors
+
+    def visit_decorator(self, dec: Decorator) -> None:
+        if dec.decorators:
+            decorators_to_store = dec.decorators.copy()
+            decorators_to_remove: list[int] = []
+            # the index of the last non-register decorator before finding a register decorator
+            # when going through decorators from top to bottom
+            last_non_register: int | None = None
+            for i, d in enumerate(decorators_to_store):
+                impl = get_singledispatch_register_call_info(d, dec.func)
+                if impl is not None:
+                    self.singledispatch_impls[impl.singledispatch_func].append(
+                        (impl.dispatch_type, dec.func)
+                    )
+                    decorators_to_remove.append(i)
+                    if last_non_register is not None:
+                        # found a register decorator after a non-register decorator, which we
+                        # don't support because we'd have to make a copy of the function before
+                        # calling the decorator so that we can call it later, which complicates
+                        # the implementation for something that is probably not commonly used
+                        self.errors.error(
+                            "Calling decorator after registering function not supported",
+                            self.current_path,
+                            decorators_to_store[last_non_register].line,
+                        )
+                else:
+                    if refers_to_fullname(d, "functools.singledispatch"):
+                        decorators_to_remove.append(i)
+                        # make sure that we still treat the function as a singledispatch function
+                        # even if we don't find any registered implementations (which might happen
+                        # if all registered implementations are registered dynamically)
+                        self.singledispatch_impls.setdefault(dec.func, [])
+                    last_non_register = i
+
+            if decorators_to_remove:
+                # calling register on a function that tries to dispatch based on type annotations
+                # raises a TypeError because compiled functions don't have an __annotations__
+                # attribute
+                self.decorators_to_remove[dec.func] = decorators_to_remove
+
+        super().visit_decorator(dec)
+
+
+class RegisteredImpl(NamedTuple):
+    singledispatch_func: FuncDef
+    dispatch_type: TypeInfo
+
+
+def get_singledispatch_register_call_info(
+    decorator: Expression, func: FuncDef
+) -> RegisteredImpl | None:
+    # @fun.register(complex)
+    # def g(arg): ...
+    if (
+        isinstance(decorator, CallExpr)
+        and len(decorator.args) == 1
+        and isinstance(decorator.args[0], RefExpr)
+    ):
+        callee = decorator.callee
+        dispatch_type = decorator.args[0].node
+        if not isinstance(dispatch_type, TypeInfo):
+            return None
+
+        if isinstance(callee, MemberExpr):
+            return registered_impl_from_possible_register_call(callee, dispatch_type)
+    # @fun.register
+    # def g(arg: int): ...
+    elif isinstance(decorator, MemberExpr):
+        # we don't know if this is a register call yet, so we can't be sure that the function
+        # actually has arguments
+        if not func.arguments:
+            return None
+        arg_type = get_proper_type(func.arguments[0].variable.type)
+        if not isinstance(arg_type, Instance):
+            return None
+        info = arg_type.type
+        return registered_impl_from_possible_register_call(decorator, info)
+    return None
+
+
+def registered_impl_from_possible_register_call(
+    expr: MemberExpr, dispatch_type: TypeInfo
+) -> RegisteredImpl | None:
+    if expr.name == "register" and isinstance(expr.expr, NameExpr):
+        node = expr.expr.node
+        if isinstance(node, Decorator):
+            return RegisteredImpl(node.func, dispatch_type)
+    return None

venv/lib/python3.12/site-packages/mypyc/irbuild/specialize.py

@@ -0,0 +1,822 @@
+"""Special case IR generation of calls to specific builtin functions.
+
+Most special cases should be handled using the data driven "primitive
+ops" system, but certain operations require special handling that has
+access to the AST/IR directly and can make decisions/optimizations
+based on it. These special cases can be implemented here.
+
+For example, we use specializers to statically emit the length of a
+fixed length tuple and to emit optimized code for any()/all() calls with
+generator comprehensions as the argument.
+
+See comment below for more documentation.
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Optional
+
+from mypy.nodes import (
+    ARG_NAMED,
+    ARG_POS,
+    CallExpr,
+    DictExpr,
+    Expression,
+    GeneratorExpr,
+    IntExpr,
+    ListExpr,
+    MemberExpr,
+    NameExpr,
+    RefExpr,
+    StrExpr,
+    TupleExpr,
+)
+from mypy.types import AnyType, TypeOfAny
+from mypyc.ir.ops import (
+    BasicBlock,
+    Extend,
+    Integer,
+    RaiseStandardError,
+    Register,
+    Truncate,
+    Unreachable,
+    Value,
+)
+from mypyc.ir.rtypes import (
+    RInstance,
+    RPrimitive,
+    RTuple,
+    RType,
+    bool_rprimitive,
+    c_int_rprimitive,
+    dict_rprimitive,
+    int16_rprimitive,
+    int32_rprimitive,
+    int64_rprimitive,
+    int_rprimitive,
+    is_bool_rprimitive,
+    is_dict_rprimitive,
+    is_fixed_width_rtype,
+    is_float_rprimitive,
+    is_int16_rprimitive,
+    is_int32_rprimitive,
+    is_int64_rprimitive,
+    is_int_rprimitive,
+    is_list_rprimitive,
+    is_uint8_rprimitive,
+    list_rprimitive,
+    set_rprimitive,
+    str_rprimitive,
+    uint8_rprimitive,
+)
+from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.for_helpers import (
+    comprehension_helper,
+    sequence_from_generator_preallocate_helper,
+    translate_list_comprehension,
+    translate_set_comprehension,
+)
+from mypyc.irbuild.format_str_tokenizer import (
+    FormatOp,
+    convert_format_expr_to_str,
+    join_formatted_strings,
+    tokenizer_format_call,
+)
+from mypyc.primitives.dict_ops import (
+    dict_items_op,
+    dict_keys_op,
+    dict_setdefault_spec_init_op,
+    dict_values_op,
+)
+from mypyc.primitives.list_ops import new_list_set_item_op
+from mypyc.primitives.tuple_ops import new_tuple_set_item_op
+
+# Specializers are attempted before compiling the arguments to the
+# function.  Specializers can return None to indicate that they failed
+# and the call should be compiled normally. Otherwise they should emit
+# code for the call and return a Value containing the result.
+#
+# Specializers take three arguments: the IRBuilder, the CallExpr being
+# compiled, and the RefExpr that is the left hand side of the call.
+Specializer = Callable[["IRBuilder", CallExpr, RefExpr], Optional[Value]]
+
+# Dictionary containing all configured specializers.
+#
+# Specializers can operate on methods as well, and are keyed on the
+# name and RType in that case.
+specializers: dict[tuple[str, RType | None], list[Specializer]] = {}
+
+
+def _apply_specialization(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr, name: str | None, typ: RType | None = None
+) -> Value | None:
+    # TODO: Allow special cases to have default args or named args. Currently they don't since
+    #       they check that everything in arg_kinds is ARG_POS.
+
+    # If there is a specializer for this function, try calling it.
+    # Return the first successful one.
+    if name and (name, typ) in specializers:
+        for specializer in specializers[name, typ]:
+            val = specializer(builder, expr, callee)
+            if val is not None:
+                return val
+    return None
+
+
+def apply_function_specialization(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr
+) -> Value | None:
+    """Invoke the Specializer callback for a function if one has been registered"""
+    return _apply_specialization(builder, expr, callee, callee.fullname)
+
+
+def apply_method_specialization(
+    builder: IRBuilder, expr: CallExpr, callee: MemberExpr, typ: RType | None = None
+) -> Value | None:
+    """Invoke the Specializer callback for a method if one has been registered"""
+    name = callee.fullname if typ is None else callee.name
+    return _apply_specialization(builder, expr, callee, name, typ)
+
+
+def specialize_function(
+    name: str, typ: RType | None = None
+) -> Callable[[Specializer], Specializer]:
+    """Decorator to register a function as being a specializer.
+
+    There may exist multiple specializers for one function. When
+    translating method calls, the earlier appended specializer has
+    higher priority.
+    """
+
+    def wrapper(f: Specializer) -> Specializer:
+        specializers.setdefault((name, typ), []).append(f)
+        return f
+
+    return wrapper
+
+
+@specialize_function("builtins.globals")
+def translate_globals(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) == 0:
+        return builder.load_globals_dict()
+    return None
+
+
+@specialize_function("builtins.abs")
+@specialize_function("builtins.int")
+@specialize_function("builtins.float")
+@specialize_function("builtins.complex")
+@specialize_function("mypy_extensions.i64")
+@specialize_function("mypy_extensions.i32")
+@specialize_function("mypy_extensions.i16")
+@specialize_function("mypy_extensions.u8")
+def translate_builtins_with_unary_dunder(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr
+) -> Value | None:
+    """Specialize calls on native classes that implement the associated dunder.
+
+    E.g. i64(x) gets specialized to x.__int__() if x is a native instance.
+    """
+    if len(expr.args) == 1 and expr.arg_kinds == [ARG_POS] and isinstance(callee, NameExpr):
+        arg = expr.args[0]
+        arg_typ = builder.node_type(arg)
+        shortname = callee.fullname.split(".")[1]
+        if shortname in ("i64", "i32", "i16", "u8"):
+            method = "__int__"
+        else:
+            method = f"__{shortname}__"
+        if isinstance(arg_typ, RInstance) and arg_typ.class_ir.has_method(method):
+            obj = builder.accept(arg)
+            return builder.gen_method_call(obj, method, [], None, expr.line)
+
+    return None
+
+
+@specialize_function("builtins.len")
+def translate_len(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) == 1 and expr.arg_kinds == [ARG_POS]:
+        arg = expr.args[0]
+        expr_rtype = builder.node_type(arg)
+        if isinstance(expr_rtype, RTuple):
+            # len() of fixed-length tuple can be trivially determined
+            # statically, though we still need to evaluate it.
+            builder.accept(arg)
+            return Integer(len(expr_rtype.types))
+        else:
+            if is_list_rprimitive(builder.node_type(arg)):
+                borrow = True
+            else:
+                borrow = False
+            obj = builder.accept(arg, can_borrow=borrow)
+            return builder.builtin_len(obj, expr.line)
+    return None
+
+
+@specialize_function("builtins.list")
+def dict_methods_fast_path(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    """Specialize a common case when list() is called on a dictionary
+    view method call.
+
+    For example:
+        foo = list(bar.keys())
+    """
+    if not (len(expr.args) == 1 and expr.arg_kinds == [ARG_POS]):
+        return None
+    arg = expr.args[0]
+    if not (isinstance(arg, CallExpr) and not arg.args and isinstance(arg.callee, MemberExpr)):
+        return None
+    base = arg.callee.expr
+    attr = arg.callee.name
+    rtype = builder.node_type(base)
+    if not (is_dict_rprimitive(rtype) and attr in ("keys", "values", "items")):
+        return None
+
+    obj = builder.accept(base)
+    # Note that it is not safe to use fast methods on dict subclasses,
+    # so the corresponding helpers in CPy.h fallback to (inlined)
+    # generic logic.
+    if attr == "keys":
+        return builder.call_c(dict_keys_op, [obj], expr.line)
+    elif attr == "values":
+        return builder.call_c(dict_values_op, [obj], expr.line)
+    else:
+        return builder.call_c(dict_items_op, [obj], expr.line)
+
+
+@specialize_function("builtins.list")
+def translate_list_from_generator_call(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr
+) -> Value | None:
+    """Special case for simplest list comprehension.
+
+    For example:
+        list(f(x) for x in some_list/some_tuple/some_str)
+    'translate_list_comprehension()' would take care of other cases
+    if this fails.
+    """
+    if (
+        len(expr.args) == 1
+        and expr.arg_kinds[0] == ARG_POS
+        and isinstance(expr.args[0], GeneratorExpr)
+    ):
+        return sequence_from_generator_preallocate_helper(
+            builder,
+            expr.args[0],
+            empty_op_llbuilder=builder.builder.new_list_op_with_length,
+            set_item_op=new_list_set_item_op,
+        )
+    return None
+
+
+@specialize_function("builtins.tuple")
+def translate_tuple_from_generator_call(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr
+) -> Value | None:
+    """Special case for simplest tuple creation from a generator.
+
+    For example:
+        tuple(f(x) for x in some_list/some_tuple/some_str)
+    'translate_safe_generator_call()' would take care of other cases
+    if this fails.
+    """
+    if (
+        len(expr.args) == 1
+        and expr.arg_kinds[0] == ARG_POS
+        and isinstance(expr.args[0], GeneratorExpr)
+    ):
+        return sequence_from_generator_preallocate_helper(
+            builder,
+            expr.args[0],
+            empty_op_llbuilder=builder.builder.new_tuple_with_length,
+            set_item_op=new_tuple_set_item_op,
+        )
+    return None
+
+
+@specialize_function("builtins.set")
+def translate_set_from_generator_call(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr
+) -> Value | None:
+    """Special case for set creation from a generator.
+
+    For example:
+        set(f(...) for ... in iterator/nested_generators...)
+    """
+    if (
+        len(expr.args) == 1
+        and expr.arg_kinds[0] == ARG_POS
+        and isinstance(expr.args[0], GeneratorExpr)
+    ):
+        return translate_set_comprehension(builder, expr.args[0])
+    return None
+
+
+@specialize_function("builtins.min")
+@specialize_function("builtins.max")
+def faster_min_max(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if expr.arg_kinds == [ARG_POS, ARG_POS]:
+        x, y = builder.accept(expr.args[0]), builder.accept(expr.args[1])
+        result = Register(builder.node_type(expr))
+        # CPython evaluates arguments reversely when calling min(...) or max(...)
+        if callee.fullname == "builtins.min":
+            comparison = builder.binary_op(y, x, "<", expr.line)
+        else:
+            comparison = builder.binary_op(y, x, ">", expr.line)
+
+        true_block, false_block, next_block = BasicBlock(), BasicBlock(), BasicBlock()
+        builder.add_bool_branch(comparison, true_block, false_block)
+
+        builder.activate_block(true_block)
+        builder.assign(result, builder.coerce(y, result.type, expr.line), expr.line)
+        builder.goto(next_block)
+
+        builder.activate_block(false_block)
+        builder.assign(result, builder.coerce(x, result.type, expr.line), expr.line)
+        builder.goto(next_block)
+
+        builder.activate_block(next_block)
+        return result
+    return None
+
+
+@specialize_function("builtins.tuple")
+@specialize_function("builtins.frozenset")
+@specialize_function("builtins.dict")
+@specialize_function("builtins.min")
+@specialize_function("builtins.max")
+@specialize_function("builtins.sorted")
+@specialize_function("collections.OrderedDict")
+@specialize_function("join", str_rprimitive)
+@specialize_function("extend", list_rprimitive)
+@specialize_function("update", dict_rprimitive)
+@specialize_function("update", set_rprimitive)
+def translate_safe_generator_call(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr
+) -> Value | None:
+    """Special cases for things that consume iterators where we know we
+    can safely compile a generator into a list.
+    """
+    if (
+        len(expr.args) > 0
+        and expr.arg_kinds[0] == ARG_POS
+        and isinstance(expr.args[0], GeneratorExpr)
+    ):
+        if isinstance(callee, MemberExpr):
+            return builder.gen_method_call(
+                builder.accept(callee.expr),
+                callee.name,
+                (
+                    [translate_list_comprehension(builder, expr.args[0])]
+                    + [builder.accept(arg) for arg in expr.args[1:]]
+                ),
+                builder.node_type(expr),
+                expr.line,
+                expr.arg_kinds,
+                expr.arg_names,
+            )
+        else:
+            return builder.call_refexpr_with_args(
+                expr,
+                callee,
+                (
+                    [translate_list_comprehension(builder, expr.args[0])]
+                    + [builder.accept(arg) for arg in expr.args[1:]]
+                ),
+            )
+    return None
+
+
+@specialize_function("builtins.any")
+def translate_any_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if (
+        len(expr.args) == 1
+        and expr.arg_kinds == [ARG_POS]
+        and isinstance(expr.args[0], GeneratorExpr)
+    ):
+        return any_all_helper(builder, expr.args[0], builder.false, lambda x: x, builder.true)
+    return None
+
+
+@specialize_function("builtins.all")
+def translate_all_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if (
+        len(expr.args) == 1
+        and expr.arg_kinds == [ARG_POS]
+        and isinstance(expr.args[0], GeneratorExpr)
+    ):
+        return any_all_helper(
+            builder,
+            expr.args[0],
+            builder.true,
+            lambda x: builder.unary_op(x, "not", expr.line),
+            builder.false,
+        )
+    return None
+
+
+def any_all_helper(
+    builder: IRBuilder,
+    gen: GeneratorExpr,
+    initial_value: Callable[[], Value],
+    modify: Callable[[Value], Value],
+    new_value: Callable[[], Value],
+) -> Value:
+    retval = Register(bool_rprimitive)
+    builder.assign(retval, initial_value(), -1)
+    loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async))
+    true_block, false_block, exit_block = BasicBlock(), BasicBlock(), BasicBlock()
+
+    def gen_inner_stmts() -> None:
+        comparison = modify(builder.accept(gen.left_expr))
+        builder.add_bool_branch(comparison, true_block, false_block)
+        builder.activate_block(true_block)
+        builder.assign(retval, new_value(), -1)
+        builder.goto(exit_block)
+        builder.activate_block(false_block)
+
+    comprehension_helper(builder, loop_params, gen_inner_stmts, gen.line)
+    builder.goto_and_activate(exit_block)
+
+    return retval
+
+
+@specialize_function("builtins.sum")
+def translate_sum_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    # specialized implementation is used if:
+    # - only one or two arguments given (if not, sum() has been given invalid arguments)
+    # - first argument is a Generator (there is no benefit to optimizing the performance of eg.
+    #   sum([1, 2, 3]), so non-Generator Iterables are not handled)
+    if not (
+        len(expr.args) in (1, 2)
+        and expr.arg_kinds[0] == ARG_POS
+        and isinstance(expr.args[0], GeneratorExpr)
+    ):
+        return None
+
+    # handle 'start' argument, if given
+    if len(expr.args) == 2:
+        # ensure call to sum() was properly constructed
+        if expr.arg_kinds[1] not in (ARG_POS, ARG_NAMED):
+            return None
+        start_expr = expr.args[1]
+    else:
+        start_expr = IntExpr(0)
+
+    gen_expr = expr.args[0]
+    target_type = builder.node_type(expr)
+    retval = Register(target_type)
+    builder.assign(retval, builder.coerce(builder.accept(start_expr), target_type, -1), -1)
+
+    def gen_inner_stmts() -> None:
+        call_expr = builder.accept(gen_expr.left_expr)
+        builder.assign(retval, builder.binary_op(retval, call_expr, "+", -1), -1)
+
+    loop_params = list(
+        zip(gen_expr.indices, gen_expr.sequences, gen_expr.condlists, gen_expr.is_async)
+    )
+    comprehension_helper(builder, loop_params, gen_inner_stmts, gen_expr.line)
+
+    return retval
+
+
+@specialize_function("dataclasses.field")
+@specialize_function("attr.ib")
+@specialize_function("attr.attrib")
+@specialize_function("attr.Factory")
+def translate_dataclasses_field_call(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr
+) -> Value | None:
+    """Special case for 'dataclasses.field', 'attr.attrib', and 'attr.Factory'
+    function calls because the results of such calls are type-checked
+    by mypy using the types of the arguments to their respective
+    functions, resulting in attempted coercions by mypyc that throw a
+    runtime error.
+    """
+    builder.types[expr] = AnyType(TypeOfAny.from_error)
+    return None
+
+
+@specialize_function("builtins.next")
+def translate_next_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    """Special case for calling next() on a generator expression, an
+    idiom that shows up some in mypy.
+
+    For example, next(x for x in l if x.id == 12, None) will
+    generate code that searches l for an element where x.id == 12
+    and produce the first such object, or None if no such element
+    exists.
+    """
+    if not (
+        expr.arg_kinds in ([ARG_POS], [ARG_POS, ARG_POS])
+        and isinstance(expr.args[0], GeneratorExpr)
+    ):
+        return None
+
+    gen = expr.args[0]
+    retval = Register(builder.node_type(expr))
+    default_val = builder.accept(expr.args[1]) if len(expr.args) > 1 else None
+    exit_block = BasicBlock()
+
+    def gen_inner_stmts() -> None:
+        # next takes the first element of the generator, so if
+        # something gets produced, we are done.
+        builder.assign(retval, builder.accept(gen.left_expr), gen.left_expr.line)
+        builder.goto(exit_block)
+
+    loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async))
+    comprehension_helper(builder, loop_params, gen_inner_stmts, gen.line)
+
+    # Now we need the case for when nothing got hit. If there was
+    # a default value, we produce it, and otherwise we raise
+    # StopIteration.
+    if default_val:
+        builder.assign(retval, default_val, gen.left_expr.line)
+        builder.goto(exit_block)
+    else:
+        builder.add(RaiseStandardError(RaiseStandardError.STOP_ITERATION, None, expr.line))
+        builder.add(Unreachable())
+
+    builder.activate_block(exit_block)
+    return retval
+
+
+@specialize_function("builtins.isinstance")
+def translate_isinstance(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    """Special case for builtins.isinstance.
+
+    Prevent coercions on the thing we are checking the instance of -
+    there is no need to coerce something to a new type before checking
+    what type it is, and the coercion could lead to bugs.
+    """
+    if (
+        len(expr.args) == 2
+        and expr.arg_kinds == [ARG_POS, ARG_POS]
+        and isinstance(expr.args[1], (RefExpr, TupleExpr))
+    ):
+        builder.types[expr.args[0]] = AnyType(TypeOfAny.from_error)
+
+        irs = builder.flatten_classes(expr.args[1])
+        if irs is not None:
+            can_borrow = all(
+                ir.is_ext_class and not ir.inherits_python and not ir.allow_interpreted_subclasses
+                for ir in irs
+            )
+            obj = builder.accept(expr.args[0], can_borrow=can_borrow)
+            return builder.builder.isinstance_helper(obj, irs, expr.line)
+    return None
+
+
+@specialize_function("setdefault", dict_rprimitive)
+def translate_dict_setdefault(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    """Special case for 'dict.setdefault' which would only construct
+    default empty collection when needed.
+
+    The dict_setdefault_spec_init_op checks whether the dict contains
+    the key and would construct the empty collection only once.
+
+    For example, this specializer works for the following cases:
+         d.setdefault(key, set()).add(value)
+         d.setdefault(key, []).append(value)
+         d.setdefault(key, {})[inner_key] = inner_val
+    """
+    if (
+        len(expr.args) == 2
+        and expr.arg_kinds == [ARG_POS, ARG_POS]
+        and isinstance(callee, MemberExpr)
+    ):
+        arg = expr.args[1]
+        if isinstance(arg, ListExpr):
+            if len(arg.items):
+                return None
+            data_type = Integer(1, c_int_rprimitive, expr.line)
+        elif isinstance(arg, DictExpr):
+            if len(arg.items):
+                return None
+            data_type = Integer(2, c_int_rprimitive, expr.line)
+        elif (
+            isinstance(arg, CallExpr)
+            and isinstance(arg.callee, NameExpr)
+            and arg.callee.fullname == "builtins.set"
+        ):
+            if len(arg.args):
+                return None
+            data_type = Integer(3, c_int_rprimitive, expr.line)
+        else:
+            return None
+
+        callee_dict = builder.accept(callee.expr)
+        key_val = builder.accept(expr.args[0])
+        return builder.call_c(
+            dict_setdefault_spec_init_op, [callee_dict, key_val, data_type], expr.line
+        )
+    return None
+
+
+@specialize_function("format", str_rprimitive)
+def translate_str_format(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if (
+        isinstance(callee, MemberExpr)
+        and isinstance(callee.expr, StrExpr)
+        and expr.arg_kinds.count(ARG_POS) == len(expr.arg_kinds)
+    ):
+        format_str = callee.expr.value
+        tokens = tokenizer_format_call(format_str)
+        if tokens is None:
+            return None
+        literals, format_ops = tokens
+        # Convert variables to strings
+        substitutions = convert_format_expr_to_str(builder, format_ops, expr.args, expr.line)
+        if substitutions is None:
+            return None
+        return join_formatted_strings(builder, literals, substitutions, expr.line)
+    return None
+
+
+@specialize_function("join", str_rprimitive)
+def translate_fstring(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    """Special case for f-string, which is translated into str.join()
+    in mypy AST.
+
+    This specializer optimizes simplest f-strings which don't contain
+    any format operation.
+    """
+    if (
+        isinstance(callee, MemberExpr)
+        and isinstance(callee.expr, StrExpr)
+        and callee.expr.value == ""
+        and expr.arg_kinds == [ARG_POS]
+        and isinstance(expr.args[0], ListExpr)
+    ):
+        for item in expr.args[0].items:
+            if isinstance(item, StrExpr):
+                continue
+            elif isinstance(item, CallExpr):
+                if not isinstance(item.callee, MemberExpr) or item.callee.name != "format":
+                    return None
+                elif (
+                    not isinstance(item.callee.expr, StrExpr) or item.callee.expr.value != "{:{}}"
+                ):
+                    return None
+
+                if not isinstance(item.args[1], StrExpr) or item.args[1].value != "":
+                    return None
+            else:
+                return None
+
+        format_ops = []
+        exprs: list[Expression] = []
+
+        for item in expr.args[0].items:
+            if isinstance(item, StrExpr) and item.value != "":
+                format_ops.append(FormatOp.STR)
+                exprs.append(item)
+            elif isinstance(item, CallExpr):
+                format_ops.append(FormatOp.STR)
+                exprs.append(item.args[0])
+
+        substitutions = convert_format_expr_to_str(builder, format_ops, exprs, expr.line)
+        if substitutions is None:
+            return None
+
+        return join_formatted_strings(builder, None, substitutions, expr.line)
+    return None
+
+
+@specialize_function("mypy_extensions.i64")
+def translate_i64(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
+        return None
+    arg = expr.args[0]
+    arg_type = builder.node_type(arg)
+    if is_int64_rprimitive(arg_type):
+        return builder.accept(arg)
+    elif is_int32_rprimitive(arg_type) or is_int16_rprimitive(arg_type):
+        val = builder.accept(arg)
+        return builder.add(Extend(val, int64_rprimitive, signed=True, line=expr.line))
+    elif is_uint8_rprimitive(arg_type):
+        val = builder.accept(arg)
+        return builder.add(Extend(val, int64_rprimitive, signed=False, line=expr.line))
+    elif is_int_rprimitive(arg_type) or is_bool_rprimitive(arg_type):
+        val = builder.accept(arg)
+        return builder.coerce(val, int64_rprimitive, expr.line)
+    return None
+
+
+@specialize_function("mypy_extensions.i32")
+def translate_i32(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
+        return None
+    arg = expr.args[0]
+    arg_type = builder.node_type(arg)
+    if is_int32_rprimitive(arg_type):
+        return builder.accept(arg)
+    elif is_int64_rprimitive(arg_type):
+        val = builder.accept(arg)
+        return builder.add(Truncate(val, int32_rprimitive, line=expr.line))
+    elif is_int16_rprimitive(arg_type):
+        val = builder.accept(arg)
+        return builder.add(Extend(val, int32_rprimitive, signed=True, line=expr.line))
+    elif is_uint8_rprimitive(arg_type):
+        val = builder.accept(arg)
+        return builder.add(Extend(val, int32_rprimitive, signed=False, line=expr.line))
+    elif is_int_rprimitive(arg_type) or is_bool_rprimitive(arg_type):
+        val = builder.accept(arg)
+        val = truncate_literal(val, int32_rprimitive)
+        return builder.coerce(val, int32_rprimitive, expr.line)
+    return None
+
+
+@specialize_function("mypy_extensions.i16")
+def translate_i16(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
+        return None
+    arg = expr.args[0]
+    arg_type = builder.node_type(arg)
+    if is_int16_rprimitive(arg_type):
+        return builder.accept(arg)
+    elif is_int32_rprimitive(arg_type) or is_int64_rprimitive(arg_type):
+        val = builder.accept(arg)
+        return builder.add(Truncate(val, int16_rprimitive, line=expr.line))
+    elif is_uint8_rprimitive(arg_type):
+        val = builder.accept(arg)
+        return builder.add(Extend(val, int16_rprimitive, signed=False, line=expr.line))
+    elif is_int_rprimitive(arg_type) or is_bool_rprimitive(arg_type):
+        val = builder.accept(arg)
+        val = truncate_literal(val, int16_rprimitive)
+        return builder.coerce(val, int16_rprimitive, expr.line)
+    return None
+
+
+@specialize_function("mypy_extensions.u8")
+def translate_u8(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
+        return None
+    arg = expr.args[0]
+    arg_type = builder.node_type(arg)
+    if is_uint8_rprimitive(arg_type):
+        return builder.accept(arg)
+    elif (
+        is_int16_rprimitive(arg_type)
+        or is_int32_rprimitive(arg_type)
+        or is_int64_rprimitive(arg_type)
+    ):
+        val = builder.accept(arg)
+        return builder.add(Truncate(val, uint8_rprimitive, line=expr.line))
+    elif is_int_rprimitive(arg_type) or is_bool_rprimitive(arg_type):
+        val = builder.accept(arg)
+        val = truncate_literal(val, uint8_rprimitive)
+        return builder.coerce(val, uint8_rprimitive, expr.line)
+    return None
+
+
+def truncate_literal(value: Value, rtype: RPrimitive) -> Value:
+    """If value is an integer literal value, truncate it to given native int rtype.
+
+    For example, truncate 256 into 0 if rtype is u8.
+    """
+    if not isinstance(value, Integer):
+        return value  # Not a literal, nothing to do
+    x = value.numeric_value()
+    max_unsigned = (1 << (rtype.size * 8)) - 1
+    x = x & max_unsigned
+    if rtype.is_signed and x >= (max_unsigned + 1) // 2:
+        # Adjust to make it a negative value
+        x -= max_unsigned + 1
+    return Integer(x, rtype)
+
+
+@specialize_function("builtins.int")
+def translate_int(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
+        return None
+    arg = expr.args[0]
+    arg_type = builder.node_type(arg)
+    if (
+        is_bool_rprimitive(arg_type)
+        or is_int_rprimitive(arg_type)
+        or is_fixed_width_rtype(arg_type)
+    ):
+        src = builder.accept(arg)
+        return builder.coerce(src, int_rprimitive, expr.line)
+    return None
+
+
+@specialize_function("builtins.bool")
+def translate_bool(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
+        return None
+    arg = expr.args[0]
+    src = builder.accept(arg)
+    return builder.builder.bool_value(src)
+
+
+@specialize_function("builtins.float")
+def translate_float(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
+        return None
+    arg = expr.args[0]
+    arg_type = builder.node_type(arg)
+    if is_float_rprimitive(arg_type):
+        # No-op float conversion.
+        return builder.accept(arg)
+    return None

venv/lib/python3.12/site-packages/mypyc/irbuild/statement.py

@@ -0,0 +1,989 @@
+"""Transform mypy statement ASTs to mypyc IR (Intermediate Representation).
+
+The top-level AST transformation logic is implemented in mypyc.irbuild.visitor
+and mypyc.irbuild.builder.
+
+A few statements are transformed in mypyc.irbuild.function (yield, for example).
+"""
+
+from __future__ import annotations
+
+import importlib.util
+from typing import Callable, Sequence
+
+from mypy.nodes import (
+    AssertStmt,
+    AssignmentStmt,
+    AwaitExpr,
+    Block,
+    BreakStmt,
+    ContinueStmt,
+    DelStmt,
+    Expression,
+    ExpressionStmt,
+    ForStmt,
+    IfStmt,
+    Import,
+    ImportAll,
+    ImportFrom,
+    ListExpr,
+    Lvalue,
+    MatchStmt,
+    OperatorAssignmentStmt,
+    RaiseStmt,
+    ReturnStmt,
+    StarExpr,
+    StrExpr,
+    TempNode,
+    TryStmt,
+    TupleExpr,
+    WhileStmt,
+    WithStmt,
+    YieldExpr,
+    YieldFromExpr,
+)
+from mypyc.ir.ops import (
+    NAMESPACE_MODULE,
+    NO_TRACEBACK_LINE_NO,
+    Assign,
+    BasicBlock,
+    Branch,
+    InitStatic,
+    Integer,
+    LoadAddress,
+    LoadErrorValue,
+    LoadLiteral,
+    LoadStatic,
+    MethodCall,
+    RaiseStandardError,
+    Register,
+    Return,
+    TupleGet,
+    Unreachable,
+    Value,
+)
+from mypyc.ir.rtypes import (
+    RInstance,
+    c_pyssize_t_rprimitive,
+    exc_rtuple,
+    is_tagged,
+    none_rprimitive,
+    object_pointer_rprimitive,
+    object_rprimitive,
+)
+from mypyc.irbuild.ast_helpers import is_borrow_friendly_expr, process_conditional
+from mypyc.irbuild.builder import IRBuilder, int_borrow_friendly_op
+from mypyc.irbuild.for_helpers import for_loop_helper
+from mypyc.irbuild.generator import add_raise_exception_blocks_to_generator_class
+from mypyc.irbuild.nonlocalcontrol import (
+    ExceptNonlocalControl,
+    FinallyNonlocalControl,
+    TryFinallyNonlocalControl,
+)
+from mypyc.irbuild.targets import (
+    AssignmentTarget,
+    AssignmentTargetAttr,
+    AssignmentTargetIndex,
+    AssignmentTargetRegister,
+    AssignmentTargetTuple,
+)
+from mypyc.primitives.exc_ops import (
+    error_catch_op,
+    exc_matches_op,
+    get_exc_info_op,
+    get_exc_value_op,
+    keep_propagating_op,
+    raise_exception_op,
+    reraise_exception_op,
+    restore_exc_info_op,
+)
+from mypyc.primitives.generic_ops import iter_op, next_raw_op, py_delattr_op
+from mypyc.primitives.misc_ops import (
+    check_stop_op,
+    coro_op,
+    import_from_many_op,
+    import_many_op,
+    send_op,
+    type_op,
+    yield_from_except_op,
+)
+
+from .match import MatchVisitor
+
+GenFunc = Callable[[], None]
+ValueGenFunc = Callable[[], Value]
+
+
+def transform_block(builder: IRBuilder, block: Block) -> None:
+    if not block.is_unreachable:
+        for stmt in block.body:
+            builder.accept(stmt)
+    # Raise a RuntimeError if we hit a non-empty unreachable block.
+    # Don't complain about empty unreachable blocks, since mypy inserts
+    # those after `if MYPY`.
+    elif block.body:
+        builder.add(
+            RaiseStandardError(
+                RaiseStandardError.RUNTIME_ERROR, "Reached allegedly unreachable code!", block.line
+            )
+        )
+        builder.add(Unreachable())
+
+
+def transform_expression_stmt(builder: IRBuilder, stmt: ExpressionStmt) -> None:
+    if isinstance(stmt.expr, StrExpr):
+        # Docstring. Ignore
+        return
+    # ExpressionStmts do not need to be coerced like other Expressions, so we shouldn't
+    # call builder.accept here.
+    stmt.expr.accept(builder.visitor)
+    builder.flush_keep_alives()
+
+
+def transform_return_stmt(builder: IRBuilder, stmt: ReturnStmt) -> None:
+    if stmt.expr:
+        retval = builder.accept(stmt.expr)
+    else:
+        retval = builder.builder.none()
+    retval = builder.coerce(retval, builder.ret_types[-1], stmt.line)
+    builder.nonlocal_control[-1].gen_return(builder, retval, stmt.line)
+
+
+def transform_assignment_stmt(builder: IRBuilder, stmt: AssignmentStmt) -> None:
+    lvalues = stmt.lvalues
+    assert lvalues
+    builder.disallow_class_assignments(lvalues, stmt.line)
+    first_lvalue = lvalues[0]
+    if stmt.type and isinstance(stmt.rvalue, TempNode):
+        # This is actually a variable annotation without initializer. Don't generate
+        # an assignment but we need to call get_assignment_target since it adds a
+        # name binding as a side effect.
+        builder.get_assignment_target(first_lvalue, stmt.line)
+        return
+
+    # Special case multiple assignments like 'x, y = e1, e2'.
+    if (
+        isinstance(first_lvalue, (TupleExpr, ListExpr))
+        and isinstance(stmt.rvalue, (TupleExpr, ListExpr))
+        and len(first_lvalue.items) == len(stmt.rvalue.items)
+        and all(is_simple_lvalue(item) for item in first_lvalue.items)
+        and len(lvalues) == 1
+    ):
+        temps = []
+        for right in stmt.rvalue.items:
+            rvalue_reg = builder.accept(right)
+            temp = Register(rvalue_reg.type)
+            builder.assign(temp, rvalue_reg, stmt.line)
+            temps.append(temp)
+        for left, temp in zip(first_lvalue.items, temps):
+            assignment_target = builder.get_assignment_target(left)
+            builder.assign(assignment_target, temp, stmt.line)
+        builder.flush_keep_alives()
+        return
+
+    line = stmt.rvalue.line
+    rvalue_reg = builder.accept(stmt.rvalue)
+    if builder.non_function_scope() and stmt.is_final_def:
+        builder.init_final_static(first_lvalue, rvalue_reg)
+    for lvalue in lvalues:
+        target = builder.get_assignment_target(lvalue)
+        builder.assign(target, rvalue_reg, line)
+        builder.flush_keep_alives()
+
+
+def is_simple_lvalue(expr: Expression) -> bool:
+    return not isinstance(expr, (StarExpr, ListExpr, TupleExpr))
+
+
+def transform_operator_assignment_stmt(builder: IRBuilder, stmt: OperatorAssignmentStmt) -> None:
+    """Operator assignment statement such as x += 1"""
+    builder.disallow_class_assignments([stmt.lvalue], stmt.line)
+    if (
+        is_tagged(builder.node_type(stmt.lvalue))
+        and is_tagged(builder.node_type(stmt.rvalue))
+        and stmt.op in int_borrow_friendly_op
+    ):
+        can_borrow = is_borrow_friendly_expr(builder, stmt.rvalue) and is_borrow_friendly_expr(
+            builder, stmt.lvalue
+        )
+    else:
+        can_borrow = False
+    target = builder.get_assignment_target(stmt.lvalue)
+    target_value = builder.read(target, stmt.line, can_borrow=can_borrow)
+    rreg = builder.accept(stmt.rvalue, can_borrow=can_borrow)
+    # the Python parser strips the '=' from operator assignment statements, so re-add it
+    op = stmt.op + "="
+    res = builder.binary_op(target_value, rreg, op, stmt.line)
+    # usually operator assignments are done in-place
+    # but when target doesn't support that we need to manually assign
+    builder.assign(target, res, res.line)
+    builder.flush_keep_alives()
+
+
+def import_globals_id_and_name(module_id: str, as_name: str | None) -> tuple[str, str]:
+    """Compute names for updating the globals dict with the appropriate module.
+
+    * For 'import foo.bar as baz' we add 'foo.bar' with the name 'baz'
+    * For 'import foo.bar' we add 'foo' with the name 'foo'
+
+    Typically we then ignore these entries and access things directly
+    via the module static, but we will use the globals version for
+    modules that mypy couldn't find, since it doesn't analyze module
+    references from those properly."""
+    if as_name:
+        globals_id = module_id
+        globals_name = as_name
+    else:
+        globals_id = globals_name = module_id.split(".")[0]
+
+    return globals_id, globals_name
+
+
+def transform_import(builder: IRBuilder, node: Import) -> None:
+    if node.is_mypy_only:
+        return
+
+    # Imports (not from imports!) are processed in an odd way so they can be
+    # table-driven and compact. Here's how it works:
+    #
+    # Import nodes are divided in groups (in the prebuild visitor). Each group
+    # consists of consecutive Import nodes:
+    #
+    #   import mod         <| group #1
+    #   import mod2         |
+    #
+    #   def foo() -> None:
+    #       import mod3    <- group #2 (*)
+    #
+    #   import mod4        <| group #3
+    #   import mod5         |
+    #
+    # Every time we encounter the first import of a group, build IR to call a
+    # helper function that will perform all of the group's imports in one go.
+    if not node.is_top_level:
+        # (*) Unless the import is within a function. In that case, prioritize
+        # speed over codesize when generating IR.
+        globals = builder.load_globals_dict()
+        for mod_id, as_name in node.ids:
+            builder.gen_import(mod_id, node.line)
+            globals_id, globals_name = import_globals_id_and_name(mod_id, as_name)
+            builder.gen_method_call(
+                globals,
+                "__setitem__",
+                [builder.load_str(globals_name), builder.get_module(globals_id, node.line)],
+                result_type=None,
+                line=node.line,
+            )
+        return
+
+    if node not in builder.module_import_groups:
+        return
+
+    modules = []
+    static_ptrs = []
+    # To show the right line number on failure, we have to add the traceback
+    # entry within the helper function (which is admittedly ugly). To drive
+    # this, we need the line number corresponding to each module.
+    mod_lines = []
+    for import_node in builder.module_import_groups[node]:
+        for mod_id, as_name in import_node.ids:
+            builder.imports[mod_id] = None
+            modules.append((mod_id, *import_globals_id_and_name(mod_id, as_name)))
+            mod_static = LoadStatic(object_rprimitive, mod_id, namespace=NAMESPACE_MODULE)
+            static_ptrs.append(builder.add(LoadAddress(object_pointer_rprimitive, mod_static)))
+            mod_lines.append(Integer(import_node.line, c_pyssize_t_rprimitive))
+
+    static_array_ptr = builder.builder.setup_rarray(object_pointer_rprimitive, static_ptrs)
+    import_line_ptr = builder.builder.setup_rarray(c_pyssize_t_rprimitive, mod_lines)
+    builder.call_c(
+        import_many_op,
+        [
+            builder.add(LoadLiteral(tuple(modules), object_rprimitive)),
+            static_array_ptr,
+            builder.load_globals_dict(),
+            builder.load_str(builder.module_path),
+            builder.load_str(builder.fn_info.name),
+            import_line_ptr,
+        ],
+        NO_TRACEBACK_LINE_NO,
+    )
+
+
+def transform_import_from(builder: IRBuilder, node: ImportFrom) -> None:
+    if node.is_mypy_only:
+        return
+
+    module_state = builder.graph[builder.module_name]
+    if module_state.ancestors is not None and module_state.ancestors:
+        module_package = module_state.ancestors[0]
+    elif builder.module_path.endswith("__init__.py"):
+        module_package = builder.module_name
+    else:
+        module_package = ""
+
+    id = importlib.util.resolve_name("." * node.relative + node.id, module_package)
+    builder.imports[id] = None
+
+    names = [name for name, _ in node.names]
+    as_names = [as_name or name for name, as_name in node.names]
+    names_literal = builder.add(LoadLiteral(tuple(names), object_rprimitive))
+    if as_names == names:
+        # Reuse names tuple to reduce verbosity.
+        as_names_literal = names_literal
+    else:
+        as_names_literal = builder.add(LoadLiteral(tuple(as_names), object_rprimitive))
+    # Note that we miscompile import from inside of functions here,
+    # since that case *shouldn't* load everything into the globals dict.
+    # This probably doesn't matter much and the code runs basically right.
+    module = builder.call_c(
+        import_from_many_op,
+        [builder.load_str(id), names_literal, as_names_literal, builder.load_globals_dict()],
+        node.line,
+    )
+    builder.add(InitStatic(module, id, namespace=NAMESPACE_MODULE))
+
+
+def transform_import_all(builder: IRBuilder, node: ImportAll) -> None:
+    if node.is_mypy_only:
+        return
+    builder.gen_import(node.id, node.line)
+
+
+def transform_if_stmt(builder: IRBuilder, stmt: IfStmt) -> None:
+    if_body, next = BasicBlock(), BasicBlock()
+    else_body = BasicBlock() if stmt.else_body else next
+
+    # If statements are normalized
+    assert len(stmt.expr) == 1
+
+    process_conditional(builder, stmt.expr[0], if_body, else_body)
+    builder.activate_block(if_body)
+    builder.accept(stmt.body[0])
+    builder.goto(next)
+    if stmt.else_body:
+        builder.activate_block(else_body)
+        builder.accept(stmt.else_body)
+        builder.goto(next)
+    builder.activate_block(next)
+
+
+def transform_while_stmt(builder: IRBuilder, s: WhileStmt) -> None:
+    body, next, top, else_block = BasicBlock(), BasicBlock(), BasicBlock(), BasicBlock()
+    normal_loop_exit = else_block if s.else_body is not None else next
+
+    builder.push_loop_stack(top, next)
+
+    # Split block so that we get a handle to the top of the loop.
+    builder.goto_and_activate(top)
+    process_conditional(builder, s.expr, body, normal_loop_exit)
+
+    builder.activate_block(body)
+    builder.accept(s.body)
+    # Add branch to the top at the end of the body.
+    builder.goto(top)
+
+    builder.pop_loop_stack()
+
+    if s.else_body is not None:
+        builder.activate_block(else_block)
+        builder.accept(s.else_body)
+        builder.goto(next)
+
+    builder.activate_block(next)
+
+
+def transform_for_stmt(builder: IRBuilder, s: ForStmt) -> None:
+    def body() -> None:
+        builder.accept(s.body)
+
+    def else_block() -> None:
+        assert s.else_body is not None
+        builder.accept(s.else_body)
+
+    for_loop_helper(
+        builder, s.index, s.expr, body, else_block if s.else_body else None, s.is_async, s.line
+    )
+
+
+def transform_break_stmt(builder: IRBuilder, node: BreakStmt) -> None:
+    builder.nonlocal_control[-1].gen_break(builder, node.line)
+
+
+def transform_continue_stmt(builder: IRBuilder, node: ContinueStmt) -> None:
+    builder.nonlocal_control[-1].gen_continue(builder, node.line)
+
+
+def transform_raise_stmt(builder: IRBuilder, s: RaiseStmt) -> None:
+    if s.expr is None:
+        builder.call_c(reraise_exception_op, [], NO_TRACEBACK_LINE_NO)
+        builder.add(Unreachable())
+        return
+
+    exc = builder.accept(s.expr)
+    builder.call_c(raise_exception_op, [exc], s.line)
+    builder.add(Unreachable())
+
+
+def transform_try_except(
+    builder: IRBuilder,
+    body: GenFunc,
+    handlers: Sequence[tuple[tuple[ValueGenFunc, int] | None, Expression | None, GenFunc]],
+    else_body: GenFunc | None,
+    line: int,
+) -> None:
+    """Generalized try/except/else handling that takes functions to gen the bodies.
+
+    The point of this is to also be able to support with."""
+    assert handlers, "try needs except"
+
+    except_entry, exit_block, cleanup_block = BasicBlock(), BasicBlock(), BasicBlock()
+    double_except_block = BasicBlock()
+    # If there is an else block, jump there after the try, otherwise just leave
+    else_block = BasicBlock() if else_body else exit_block
+
+    # Compile the try block with an error handler
+    builder.builder.push_error_handler(except_entry)
+    builder.goto_and_activate(BasicBlock())
+    body()
+    builder.goto(else_block)
+    builder.builder.pop_error_handler()
+
+    # The error handler catches the error and then checks it
+    # against the except clauses. We compile the error handler
+    # itself with an error handler so that it can properly restore
+    # the *old* exc_info if an exception occurs.
+    # The exception chaining will be done automatically when the
+    # exception is raised, based on the exception in exc_info.
+    builder.builder.push_error_handler(double_except_block)
+    builder.activate_block(except_entry)
+    old_exc = builder.maybe_spill(builder.call_c(error_catch_op, [], line))
+    # Compile the except blocks with the nonlocal control flow overridden to clear exc_info
+    builder.nonlocal_control.append(ExceptNonlocalControl(builder.nonlocal_control[-1], old_exc))
+
+    # Process the bodies
+    for type, var, handler_body in handlers:
+        next_block = None
+        if type:
+            type_f, type_line = type
+            next_block, body_block = BasicBlock(), BasicBlock()
+            matches = builder.call_c(exc_matches_op, [type_f()], type_line)
+            builder.add(Branch(matches, body_block, next_block, Branch.BOOL))
+            builder.activate_block(body_block)
+        if var:
+            target = builder.get_assignment_target(var)
+            builder.assign(target, builder.call_c(get_exc_value_op, [], var.line), var.line)
+        handler_body()
+        builder.goto(cleanup_block)
+        if next_block:
+            builder.activate_block(next_block)
+
+    # Reraise the exception if needed
+    if next_block:
+        builder.call_c(reraise_exception_op, [], NO_TRACEBACK_LINE_NO)
+        builder.add(Unreachable())
+
+    builder.nonlocal_control.pop()
+    builder.builder.pop_error_handler()
+
+    # Cleanup for if we leave except through normal control flow:
+    # restore the saved exc_info information and continue propagating
+    # the exception if it exists.
+    builder.activate_block(cleanup_block)
+    builder.call_c(restore_exc_info_op, [builder.read(old_exc)], line)
+    builder.goto(exit_block)
+
+    # Cleanup for if we leave except through a raised exception:
+    # restore the saved exc_info information and continue propagating
+    # the exception.
+    builder.activate_block(double_except_block)
+    builder.call_c(restore_exc_info_op, [builder.read(old_exc)], line)
+    builder.call_c(keep_propagating_op, [], NO_TRACEBACK_LINE_NO)
+    builder.add(Unreachable())
+
+    # If present, compile the else body in the obvious way
+    if else_body:
+        builder.activate_block(else_block)
+        else_body()
+        builder.goto(exit_block)
+
+    builder.activate_block(exit_block)
+
+
+def transform_try_except_stmt(builder: IRBuilder, t: TryStmt) -> None:
+    def body() -> None:
+        builder.accept(t.body)
+
+    # Work around scoping woes
+    def make_handler(body: Block) -> GenFunc:
+        return lambda: builder.accept(body)
+
+    def make_entry(type: Expression) -> tuple[ValueGenFunc, int]:
+        return (lambda: builder.accept(type), type.line)
+
+    handlers = [
+        (make_entry(type) if type else None, var, make_handler(body))
+        for type, var, body in zip(t.types, t.vars, t.handlers)
+    ]
+    else_body = (lambda: builder.accept(t.else_body)) if t.else_body else None
+    transform_try_except(builder, body, handlers, else_body, t.line)
+
+
+def try_finally_try(
+    builder: IRBuilder,
+    err_handler: BasicBlock,
+    return_entry: BasicBlock,
+    main_entry: BasicBlock,
+    try_body: GenFunc,
+) -> Register | AssignmentTarget | None:
+    # Compile the try block with an error handler
+    control = TryFinallyNonlocalControl(return_entry)
+    builder.builder.push_error_handler(err_handler)
+
+    builder.nonlocal_control.append(control)
+    builder.goto_and_activate(BasicBlock())
+    try_body()
+    builder.goto(main_entry)
+    builder.nonlocal_control.pop()
+    builder.builder.pop_error_handler()
+
+    return control.ret_reg
+
+
+def try_finally_entry_blocks(
+    builder: IRBuilder,
+    err_handler: BasicBlock,
+    return_entry: BasicBlock,
+    main_entry: BasicBlock,
+    finally_block: BasicBlock,
+    ret_reg: Register | AssignmentTarget | None,
+) -> Value:
+    old_exc = Register(exc_rtuple)
+
+    # Entry block for non-exceptional flow
+    builder.activate_block(main_entry)
+    if ret_reg:
+        builder.assign(ret_reg, builder.add(LoadErrorValue(builder.ret_types[-1])), -1)
+    builder.goto(return_entry)
+
+    builder.activate_block(return_entry)
+    builder.add(Assign(old_exc, builder.add(LoadErrorValue(exc_rtuple))))
+    builder.goto(finally_block)
+
+    # Entry block for errors
+    builder.activate_block(err_handler)
+    if ret_reg:
+        builder.assign(ret_reg, builder.add(LoadErrorValue(builder.ret_types[-1])), -1)
+    builder.add(Assign(old_exc, builder.call_c(error_catch_op, [], -1)))
+    builder.goto(finally_block)
+
+    return old_exc
+
+
+def try_finally_body(
+    builder: IRBuilder, finally_block: BasicBlock, finally_body: GenFunc, old_exc: Value
+) -> tuple[BasicBlock, FinallyNonlocalControl]:
+    cleanup_block = BasicBlock()
+    # Compile the finally block with the nonlocal control flow overridden to restore exc_info
+    builder.builder.push_error_handler(cleanup_block)
+    finally_control = FinallyNonlocalControl(builder.nonlocal_control[-1], old_exc)
+    builder.nonlocal_control.append(finally_control)
+    builder.activate_block(finally_block)
+    finally_body()
+    builder.nonlocal_control.pop()
+
+    return cleanup_block, finally_control
+
+
+def try_finally_resolve_control(
+    builder: IRBuilder,
+    cleanup_block: BasicBlock,
+    finally_control: FinallyNonlocalControl,
+    old_exc: Value,
+    ret_reg: Register | AssignmentTarget | None,
+) -> BasicBlock:
+    """Resolve the control flow out of a finally block.
+
+    This means returning if there was a return, propagating
+    exceptions, break/continue (soon), or just continuing on.
+    """
+    reraise, rest = BasicBlock(), BasicBlock()
+    builder.add(Branch(old_exc, rest, reraise, Branch.IS_ERROR))
+
+    # Reraise the exception if there was one
+    builder.activate_block(reraise)
+    builder.call_c(reraise_exception_op, [], NO_TRACEBACK_LINE_NO)
+    builder.add(Unreachable())
+    builder.builder.pop_error_handler()
+
+    # If there was a return, keep returning
+    if ret_reg:
+        builder.activate_block(rest)
+        return_block, rest = BasicBlock(), BasicBlock()
+        builder.add(Branch(builder.read(ret_reg), rest, return_block, Branch.IS_ERROR))
+
+        builder.activate_block(return_block)
+        builder.nonlocal_control[-1].gen_return(builder, builder.read(ret_reg), -1)
+
+    # TODO: handle break/continue
+    builder.activate_block(rest)
+    out_block = BasicBlock()
+    builder.goto(out_block)
+
+    # If there was an exception, restore again
+    builder.activate_block(cleanup_block)
+    finally_control.gen_cleanup(builder, -1)
+    builder.call_c(keep_propagating_op, [], NO_TRACEBACK_LINE_NO)
+    builder.add(Unreachable())
+
+    return out_block
+
+
+def transform_try_finally_stmt(
+    builder: IRBuilder, try_body: GenFunc, finally_body: GenFunc
+) -> None:
+    """Generalized try/finally handling that takes functions to gen the bodies.
+
+    The point of this is to also be able to support with."""
+    # Finally is a big pain, because there are so many ways that
+    # exits can occur. We emit 10+ basic blocks for every finally!
+
+    err_handler, main_entry, return_entry, finally_block = (
+        BasicBlock(),
+        BasicBlock(),
+        BasicBlock(),
+        BasicBlock(),
+    )
+
+    # Compile the body of the try
+    ret_reg = try_finally_try(builder, err_handler, return_entry, main_entry, try_body)
+
+    # Set up the entry blocks for the finally statement
+    old_exc = try_finally_entry_blocks(
+        builder, err_handler, return_entry, main_entry, finally_block, ret_reg
+    )
+
+    # Compile the body of the finally
+    cleanup_block, finally_control = try_finally_body(
+        builder, finally_block, finally_body, old_exc
+    )
+
+    # Resolve the control flow out of the finally block
+    out_block = try_finally_resolve_control(
+        builder, cleanup_block, finally_control, old_exc, ret_reg
+    )
+
+    builder.activate_block(out_block)
+
+
+def transform_try_stmt(builder: IRBuilder, t: TryStmt) -> None:
+    # Our compilation strategy for try/except/else/finally is to
+    # treat try/except/else and try/finally as separate language
+    # constructs that we compile separately. When we have a
+    # try/except/else/finally, we treat the try/except/else as the
+    # body of a try/finally block.
+    if t.is_star:
+        builder.error("Exception groups and except* cannot be compiled yet", t.line)
+    if t.finally_body:
+
+        def transform_try_body() -> None:
+            if t.handlers:
+                transform_try_except_stmt(builder, t)
+            else:
+                builder.accept(t.body)
+
+        body = t.finally_body
+
+        transform_try_finally_stmt(builder, transform_try_body, lambda: builder.accept(body))
+    else:
+        transform_try_except_stmt(builder, t)
+
+
+def get_sys_exc_info(builder: IRBuilder) -> list[Value]:
+    exc_info = builder.call_c(get_exc_info_op, [], -1)
+    return [builder.add(TupleGet(exc_info, i, -1)) for i in range(3)]
+
+
+def transform_with(
+    builder: IRBuilder,
+    expr: Expression,
+    target: Lvalue | None,
+    body: GenFunc,
+    is_async: bool,
+    line: int,
+) -> None:
+    # This is basically a straight transcription of the Python code in PEP 343.
+    # I don't actually understand why a bunch of it is the way it is.
+    # We could probably optimize the case where the manager is compiled by us,
+    # but that is not our common case at all, so.
+
+    al = "a" if is_async else ""
+
+    mgr_v = builder.accept(expr)
+    is_native = isinstance(mgr_v.type, RInstance)
+    if is_native:
+        value = builder.add(MethodCall(mgr_v, f"__{al}enter__", args=[], line=line))
+        exit_ = None
+    else:
+        typ = builder.call_c(type_op, [mgr_v], line)
+        exit_ = builder.maybe_spill(builder.py_get_attr(typ, f"__{al}exit__", line))
+        value = builder.py_call(builder.py_get_attr(typ, f"__{al}enter__", line), [mgr_v], line)
+
+    mgr = builder.maybe_spill(mgr_v)
+    exc = builder.maybe_spill_assignable(builder.true())
+    if is_async:
+        value = emit_await(builder, value, line)
+
+    def maybe_natively_call_exit(exc_info: bool) -> Value:
+        if exc_info:
+            args = get_sys_exc_info(builder)
+        else:
+            none = builder.none_object()
+            args = [none, none, none]
+
+        if is_native:
+            assert isinstance(mgr_v.type, RInstance)
+            exit_val = builder.gen_method_call(
+                builder.read(mgr),
+                f"__{al}exit__",
+                arg_values=args,
+                line=line,
+                result_type=none_rprimitive,
+            )
+        else:
+            assert exit_ is not None
+            exit_val = builder.py_call(builder.read(exit_), [builder.read(mgr)] + args, line)
+
+        if is_async:
+            return emit_await(builder, exit_val, line)
+        else:
+            return exit_val
+
+    def try_body() -> None:
+        if target:
+            builder.assign(builder.get_assignment_target(target), value, line)
+        body()
+
+    def except_body() -> None:
+        builder.assign(exc, builder.false(), line)
+        out_block, reraise_block = BasicBlock(), BasicBlock()
+        builder.add_bool_branch(maybe_natively_call_exit(exc_info=True), out_block, reraise_block)
+        builder.activate_block(reraise_block)
+        builder.call_c(reraise_exception_op, [], NO_TRACEBACK_LINE_NO)
+        builder.add(Unreachable())
+        builder.activate_block(out_block)
+
+    def finally_body() -> None:
+        out_block, exit_block = BasicBlock(), BasicBlock()
+        builder.add(Branch(builder.read(exc), exit_block, out_block, Branch.BOOL))
+        builder.activate_block(exit_block)
+
+        maybe_natively_call_exit(exc_info=False)
+        builder.goto_and_activate(out_block)
+
+    transform_try_finally_stmt(
+        builder,
+        lambda: transform_try_except(builder, try_body, [(None, None, except_body)], None, line),
+        finally_body,
+    )
+
+
+def transform_with_stmt(builder: IRBuilder, o: WithStmt) -> None:
+    # Generate separate logic for each expr in it, left to right
+    def generate(i: int) -> None:
+        if i >= len(o.expr):
+            builder.accept(o.body)
+        else:
+            transform_with(
+                builder, o.expr[i], o.target[i], lambda: generate(i + 1), o.is_async, o.line
+            )
+
+    generate(0)
+
+
+def transform_assert_stmt(builder: IRBuilder, a: AssertStmt) -> None:
+    if builder.options.strip_asserts:
+        return
+    cond = builder.accept(a.expr)
+    ok_block, error_block = BasicBlock(), BasicBlock()
+    builder.add_bool_branch(cond, ok_block, error_block)
+    builder.activate_block(error_block)
+    if a.msg is None:
+        # Special case (for simpler generated code)
+        builder.add(RaiseStandardError(RaiseStandardError.ASSERTION_ERROR, None, a.line))
+    elif isinstance(a.msg, StrExpr):
+        # Another special case
+        builder.add(RaiseStandardError(RaiseStandardError.ASSERTION_ERROR, a.msg.value, a.line))
+    else:
+        # The general case -- explicitly construct an exception instance
+        message = builder.accept(a.msg)
+        exc_type = builder.load_module_attr_by_fullname("builtins.AssertionError", a.line)
+        exc = builder.py_call(exc_type, [message], a.line)
+        builder.call_c(raise_exception_op, [exc], a.line)
+    builder.add(Unreachable())
+    builder.activate_block(ok_block)
+
+
+def transform_del_stmt(builder: IRBuilder, o: DelStmt) -> None:
+    transform_del_item(builder, builder.get_assignment_target(o.expr), o.line)
+
+
+def transform_del_item(builder: IRBuilder, target: AssignmentTarget, line: int) -> None:
+    if isinstance(target, AssignmentTargetIndex):
+        builder.gen_method_call(
+            target.base, "__delitem__", [target.index], result_type=None, line=line
+        )
+    elif isinstance(target, AssignmentTargetAttr):
+        if isinstance(target.obj_type, RInstance):
+            cl = target.obj_type.class_ir
+            if not cl.is_deletable(target.attr):
+                builder.error(f'"{target.attr}" cannot be deleted', line)
+                builder.note(
+                    'Using "__deletable__ = '
+                    + '[\'<attr>\']" in the class body enables "del obj.<attr>"',
+                    line,
+                )
+        key = builder.load_str(target.attr)
+        builder.call_c(py_delattr_op, [target.obj, key], line)
+    elif isinstance(target, AssignmentTargetRegister):
+        # Delete a local by assigning an error value to it, which will
+        # prompt the insertion of uninit checks.
+        builder.add(
+            Assign(target.register, builder.add(LoadErrorValue(target.type, undefines=True)))
+        )
+    elif isinstance(target, AssignmentTargetTuple):
+        for subtarget in target.items:
+            transform_del_item(builder, subtarget, line)
+
+
+# yield/yield from/await
+
+# These are really expressions, not statements... but they depend on try/except/finally
+
+
+def emit_yield(builder: IRBuilder, val: Value, line: int) -> Value:
+    retval = builder.coerce(val, builder.ret_types[-1], line)
+
+    cls = builder.fn_info.generator_class
+    # Create a new block for the instructions immediately following the yield expression, and
+    # set the next label so that the next time '__next__' is called on the generator object,
+    # the function continues at the new block.
+    next_block = BasicBlock()
+    next_label = len(cls.continuation_blocks)
+    cls.continuation_blocks.append(next_block)
+    builder.assign(cls.next_label_target, Integer(next_label), line)
+    builder.add(Return(retval))
+    builder.activate_block(next_block)
+
+    add_raise_exception_blocks_to_generator_class(builder, line)
+
+    assert cls.send_arg_reg is not None
+    return cls.send_arg_reg
+
+
+def emit_yield_from_or_await(
+    builder: IRBuilder, val: Value, line: int, *, is_await: bool
+) -> Value:
+    # This is basically an implementation of the code in PEP 380.
+
+    # TODO: do we want to use the right types here?
+    result = Register(object_rprimitive)
+    to_yield_reg = Register(object_rprimitive)
+    received_reg = Register(object_rprimitive)
+
+    get_op = coro_op if is_await else iter_op
+    iter_val = builder.call_c(get_op, [val], line)
+
+    iter_reg = builder.maybe_spill_assignable(iter_val)
+
+    stop_block, main_block, done_block = BasicBlock(), BasicBlock(), BasicBlock()
+    _y_init = builder.call_c(next_raw_op, [builder.read(iter_reg)], line)
+    builder.add(Branch(_y_init, stop_block, main_block, Branch.IS_ERROR))
+
+    # Try extracting a return value from a StopIteration and return it.
+    # If it wasn't, this reraises the exception.
+    builder.activate_block(stop_block)
+    builder.assign(result, builder.call_c(check_stop_op, [], line), line)
+    builder.goto(done_block)
+
+    builder.activate_block(main_block)
+    builder.assign(to_yield_reg, _y_init, line)
+
+    # OK Now the main loop!
+    loop_block = BasicBlock()
+    builder.goto_and_activate(loop_block)
+
+    def try_body() -> None:
+        builder.assign(received_reg, emit_yield(builder, builder.read(to_yield_reg), line), line)
+
+    def except_body() -> None:
+        # The body of the except is all implemented in a C function to
+        # reduce how much code we need to generate. It returns a value
+        # indicating whether to break or yield (or raise an exception).
+        val = Register(object_rprimitive)
+        val_address = builder.add(LoadAddress(object_pointer_rprimitive, val))
+        to_stop = builder.call_c(yield_from_except_op, [builder.read(iter_reg), val_address], line)
+
+        ok, stop = BasicBlock(), BasicBlock()
+        builder.add(Branch(to_stop, stop, ok, Branch.BOOL))
+
+        # The exception got swallowed. Continue, yielding the returned value
+        builder.activate_block(ok)
+        builder.assign(to_yield_reg, val, line)
+        builder.nonlocal_control[-1].gen_continue(builder, line)
+
+        # The exception was a StopIteration. Stop iterating.
+        builder.activate_block(stop)
+        builder.assign(result, val, line)
+        builder.nonlocal_control[-1].gen_break(builder, line)
+
+    def else_body() -> None:
+        # Do a next() or a .send(). It will return NULL on exception
+        # but it won't automatically propagate.
+        _y = builder.call_c(send_op, [builder.read(iter_reg), builder.read(received_reg)], line)
+        ok, stop = BasicBlock(), BasicBlock()
+        builder.add(Branch(_y, stop, ok, Branch.IS_ERROR))
+
+        # Everything's fine. Yield it.
+        builder.activate_block(ok)
+        builder.assign(to_yield_reg, _y, line)
+        builder.nonlocal_control[-1].gen_continue(builder, line)
+
+        # Try extracting a return value from a StopIteration and return it.
+        # If it wasn't, this rereaises the exception.
+        builder.activate_block(stop)
+        builder.assign(result, builder.call_c(check_stop_op, [], line), line)
+        builder.nonlocal_control[-1].gen_break(builder, line)
+
+    builder.push_loop_stack(loop_block, done_block)
+    transform_try_except(builder, try_body, [(None, None, except_body)], else_body, line)
+    builder.pop_loop_stack()
+
+    builder.goto_and_activate(done_block)
+    return builder.read(result)
+
+
+def emit_await(builder: IRBuilder, val: Value, line: int) -> Value:
+    return emit_yield_from_or_await(builder, val, line, is_await=True)
+
+
+def transform_yield_expr(builder: IRBuilder, expr: YieldExpr) -> Value:
+    if builder.fn_info.is_coroutine:
+        builder.error("async generators are unimplemented", expr.line)
+
+    if expr.expr:
+        retval = builder.accept(expr.expr)
+    else:
+        retval = builder.builder.none()
+    return emit_yield(builder, retval, expr.line)
+
+
+def transform_yield_from_expr(builder: IRBuilder, o: YieldFromExpr) -> Value:
+    return emit_yield_from_or_await(builder, builder.accept(o.expr), o.line, is_await=False)
+
+
+def transform_await_expr(builder: IRBuilder, o: AwaitExpr) -> Value:
+    return emit_yield_from_or_await(builder, builder.accept(o.expr), o.line, is_await=True)
+
+
+def transform_match_stmt(builder: IRBuilder, m: MatchStmt) -> None:
+    m.accept(MatchVisitor(builder, m))

venv/lib/python3.12/site-packages/mypyc/irbuild/targets.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from mypyc.ir.ops import Register, Value
+from mypyc.ir.rtypes import RInstance, RType, object_rprimitive
+
+
+class AssignmentTarget:
+    """Abstract base class for assignment targets during IR building."""
+
+    type: RType = object_rprimitive
+
+
+class AssignmentTargetRegister(AssignmentTarget):
+    """Register as an assignment target.
+
+    This is used for local variables and some temporaries.
+    """
+
+    def __init__(self, register: Register) -> None:
+        self.register = register
+        self.type = register.type
+
+
+class AssignmentTargetIndex(AssignmentTarget):
+    """base[index] as assignment target"""
+
+    def __init__(self, base: Value, index: Value) -> None:
+        self.base = base
+        self.index = index
+        # TODO: object_rprimitive won't be right for user-defined classes. Store the
+        #       lvalue type in mypy and use a better type to avoid unneeded boxing.
+        self.type = object_rprimitive
+
+
+class AssignmentTargetAttr(AssignmentTarget):
+    """obj.attr as assignment target"""
+
+    def __init__(self, obj: Value, attr: str, can_borrow: bool = False) -> None:
+        self.obj = obj
+        self.attr = attr
+        self.can_borrow = can_borrow
+        if isinstance(obj.type, RInstance) and obj.type.class_ir.has_attr(attr):
+            # Native attribute reference
+            self.obj_type: RType = obj.type
+            self.type = obj.type.attr_type(attr)
+        else:
+            # Python attribute reference
+            self.obj_type = object_rprimitive
+            self.type = object_rprimitive
+
+
+class AssignmentTargetTuple(AssignmentTarget):
+    """x, ..., y as assignment target"""
+
+    def __init__(self, items: list[AssignmentTarget], star_idx: int | None = None) -> None:
+        self.items = items
+        self.star_idx = star_idx

venv/lib/python3.12/site-packages/mypyc/irbuild/util.py

@@ -0,0 +1,189 @@
+"""Various utilities that don't depend on other modules in mypyc.irbuild."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from mypy.nodes import (
+    ARG_NAMED,
+    ARG_NAMED_OPT,
+    ARG_OPT,
+    ARG_POS,
+    GDEF,
+    ArgKind,
+    BytesExpr,
+    CallExpr,
+    ClassDef,
+    Decorator,
+    Expression,
+    FloatExpr,
+    FuncDef,
+    IntExpr,
+    NameExpr,
+    OverloadedFuncDef,
+    RefExpr,
+    StrExpr,
+    TupleExpr,
+    UnaryExpr,
+    Var,
+)
+
+DATACLASS_DECORATORS = {"dataclasses.dataclass", "attr.s", "attr.attrs"}
+
+
+def is_trait_decorator(d: Expression) -> bool:
+    return isinstance(d, RefExpr) and d.fullname == "mypy_extensions.trait"
+
+
+def is_trait(cdef: ClassDef) -> bool:
+    return any(is_trait_decorator(d) for d in cdef.decorators) or cdef.info.is_protocol
+
+
+def dataclass_decorator_type(d: Expression) -> str | None:
+    if isinstance(d, RefExpr) and d.fullname in DATACLASS_DECORATORS:
+        return d.fullname.split(".")[0]
+    elif (
+        isinstance(d, CallExpr)
+        and isinstance(d.callee, RefExpr)
+        and d.callee.fullname in DATACLASS_DECORATORS
+    ):
+        name = d.callee.fullname.split(".")[0]
+        if name == "attr" and "auto_attribs" in d.arg_names:
+            # Note: the mypy attrs plugin checks that the value of auto_attribs is
+            # not computed at runtime, so we don't need to perform that check here
+            auto = d.args[d.arg_names.index("auto_attribs")]
+            if isinstance(auto, NameExpr) and auto.name == "True":
+                return "attr-auto"
+        return name
+    else:
+        return None
+
+
+def is_dataclass_decorator(d: Expression) -> bool:
+    return dataclass_decorator_type(d) is not None
+
+
+def is_dataclass(cdef: ClassDef) -> bool:
+    return any(is_dataclass_decorator(d) for d in cdef.decorators)
+
+
+def dataclass_type(cdef: ClassDef) -> str | None:
+    for d in cdef.decorators:
+        typ = dataclass_decorator_type(d)
+        if typ is not None:
+            return typ
+    return None
+
+
+def get_mypyc_attr_literal(e: Expression) -> Any:
+    """Convert an expression from a mypyc_attr decorator to a value.
+
+    Supports a pretty limited range."""
+    if isinstance(e, (StrExpr, IntExpr, FloatExpr)):
+        return e.value
+    elif isinstance(e, RefExpr) and e.fullname == "builtins.True":
+        return True
+    elif isinstance(e, RefExpr) and e.fullname == "builtins.False":
+        return False
+    elif isinstance(e, RefExpr) and e.fullname == "builtins.None":
+        return None
+    return NotImplemented
+
+
+def get_mypyc_attr_call(d: Expression) -> CallExpr | None:
+    """Check if an expression is a call to mypyc_attr and return it if so."""
+    if (
+        isinstance(d, CallExpr)
+        and isinstance(d.callee, RefExpr)
+        and d.callee.fullname == "mypy_extensions.mypyc_attr"
+    ):
+        return d
+    return None
+
+
+def get_mypyc_attrs(stmt: ClassDef | Decorator) -> dict[str, Any]:
+    """Collect all the mypyc_attr attributes on a class definition or a function."""
+    attrs: dict[str, Any] = {}
+    for dec in stmt.decorators:
+        d = get_mypyc_attr_call(dec)
+        if d:
+            for name, arg in zip(d.arg_names, d.args):
+                if name is None:
+                    if isinstance(arg, StrExpr):
+                        attrs[arg.value] = True
+                else:
+                    attrs[name] = get_mypyc_attr_literal(arg)
+
+    return attrs
+
+
+def is_extension_class(cdef: ClassDef) -> bool:
+    if any(
+        not is_trait_decorator(d) and not is_dataclass_decorator(d) and not get_mypyc_attr_call(d)
+        for d in cdef.decorators
+    ):
+        return False
+    if cdef.info.typeddict_type:
+        return False
+    if cdef.info.is_named_tuple:
+        return False
+    if cdef.info.metaclass_type and cdef.info.metaclass_type.type.fullname not in (
+        "abc.ABCMeta",
+        "typing.TypingMeta",
+        "typing.GenericMeta",
+    ):
+        return False
+    return True
+
+
+def get_func_def(op: FuncDef | Decorator | OverloadedFuncDef) -> FuncDef:
+    if isinstance(op, OverloadedFuncDef):
+        assert op.impl
+        op = op.impl
+    if isinstance(op, Decorator):
+        op = op.func
+    return op
+
+
+def concrete_arg_kind(kind: ArgKind) -> ArgKind:
+    """Find the concrete version of an arg kind that is being passed."""
+    if kind == ARG_OPT:
+        return ARG_POS
+    elif kind == ARG_NAMED_OPT:
+        return ARG_NAMED
+    else:
+        return kind
+
+
+def is_constant(e: Expression) -> bool:
+    """Check whether we allow an expression to appear as a default value.
+
+    We don't currently properly support storing the evaluated
+    values for default arguments and default attribute values, so
+    we restrict what expressions we allow.  We allow literals of
+    primitives types, None, and references to Final global
+    variables.
+    """
+    return (
+        isinstance(e, (StrExpr, BytesExpr, IntExpr, FloatExpr))
+        or (isinstance(e, UnaryExpr) and e.op == "-" and isinstance(e.expr, (IntExpr, FloatExpr)))
+        or (isinstance(e, TupleExpr) and all(is_constant(e) for e in e.items))
+        or (
+            isinstance(e, RefExpr)
+            and e.kind == GDEF
+            and (
+                e.fullname in ("builtins.True", "builtins.False", "builtins.None")
+                or (isinstance(e.node, Var) and e.node.is_final)
+            )
+        )
+    )
+
+
+def bytes_from_str(value: str) -> bytes:
+    """Convert a string representing bytes into actual bytes.
+
+    This is needed because the literal characters of BytesExpr (the
+    characters inside b'') are stored in BytesExpr.value, whose type is
+    'str' not 'bytes'.
+    """
+    return bytes(value, "utf8").decode("unicode-escape").encode("raw-unicode-escape")

venv/lib/python3.12/site-packages/mypyc/irbuild/visitor.py

@@ -0,0 +1,397 @@
+"""Dispatcher used when transforming a mypy AST to the IR form.
+
+mypyc.irbuild.builder and mypyc.irbuild.main are closely related.
+"""
+
+from __future__ import annotations
+
+from typing import NoReturn
+
+from mypy.nodes import (
+    AssertStmt,
+    AssertTypeExpr,
+    AssignmentExpr,
+    AssignmentStmt,
+    AwaitExpr,
+    Block,
+    BreakStmt,
+    BytesExpr,
+    CallExpr,
+    CastExpr,
+    ClassDef,
+    ComparisonExpr,
+    ComplexExpr,
+    ConditionalExpr,
+    ContinueStmt,
+    Decorator,
+    DelStmt,
+    DictExpr,
+    DictionaryComprehension,
+    EllipsisExpr,
+    EnumCallExpr,
+    ExpressionStmt,
+    FloatExpr,
+    ForStmt,
+    FuncDef,
+    GeneratorExpr,
+    GlobalDecl,
+    IfStmt,
+    Import,
+    ImportAll,
+    ImportFrom,
+    IndexExpr,
+    IntExpr,
+    LambdaExpr,
+    ListComprehension,
+    ListExpr,
+    MatchStmt,
+    MemberExpr,
+    MypyFile,
+    NamedTupleExpr,
+    NameExpr,
+    NewTypeExpr,
+    NonlocalDecl,
+    OperatorAssignmentStmt,
+    OpExpr,
+    OverloadedFuncDef,
+    ParamSpecExpr,
+    PassStmt,
+    PromoteExpr,
+    RaiseStmt,
+    ReturnStmt,
+    RevealExpr,
+    SetComprehension,
+    SetExpr,
+    SliceExpr,
+    StarExpr,
+    StrExpr,
+    SuperExpr,
+    TempNode,
+    TryStmt,
+    TupleExpr,
+    TypeAliasExpr,
+    TypeApplication,
+    TypedDictExpr,
+    TypeVarExpr,
+    TypeVarTupleExpr,
+    UnaryExpr,
+    Var,
+    WhileStmt,
+    WithStmt,
+    YieldExpr,
+    YieldFromExpr,
+)
+from mypyc.ir.ops import Value
+from mypyc.irbuild.builder import IRBuilder, IRVisitor, UnsupportedException
+from mypyc.irbuild.classdef import transform_class_def
+from mypyc.irbuild.expression import (
+    transform_assignment_expr,
+    transform_bytes_expr,
+    transform_call_expr,
+    transform_comparison_expr,
+    transform_complex_expr,
+    transform_conditional_expr,
+    transform_dict_expr,
+    transform_dictionary_comprehension,
+    transform_ellipsis,
+    transform_float_expr,
+    transform_generator_expr,
+    transform_index_expr,
+    transform_int_expr,
+    transform_list_comprehension,
+    transform_list_expr,
+    transform_member_expr,
+    transform_name_expr,
+    transform_op_expr,
+    transform_set_comprehension,
+    transform_set_expr,
+    transform_slice_expr,
+    transform_str_expr,
+    transform_super_expr,
+    transform_tuple_expr,
+    transform_unary_expr,
+)
+from mypyc.irbuild.function import (
+    transform_decorator,
+    transform_func_def,
+    transform_lambda_expr,
+    transform_overloaded_func_def,
+)
+from mypyc.irbuild.statement import (
+    transform_assert_stmt,
+    transform_assignment_stmt,
+    transform_await_expr,
+    transform_block,
+    transform_break_stmt,
+    transform_continue_stmt,
+    transform_del_stmt,
+    transform_expression_stmt,
+    transform_for_stmt,
+    transform_if_stmt,
+    transform_import,
+    transform_import_all,
+    transform_import_from,
+    transform_match_stmt,
+    transform_operator_assignment_stmt,
+    transform_raise_stmt,
+    transform_return_stmt,
+    transform_try_stmt,
+    transform_while_stmt,
+    transform_with_stmt,
+    transform_yield_expr,
+    transform_yield_from_expr,
+)
+
+
+class IRBuilderVisitor(IRVisitor):
+    """Mypy node visitor that dispatches to node transform implementations.
+
+    This class should have no non-trivial logic.
+
+    This visitor is separated from the rest of code to improve modularity and
+    to avoid import cycles.
+
+    This is based on the visitor pattern
+    (https://en.wikipedia.org/wiki/Visitor_pattern).
+    """
+
+    # This gets passed to all the implementations and contains all the
+    # state and many helpers. The attribute is initialized outside
+    # this class since this class and IRBuilder form a reference loop.
+    builder: IRBuilder
+
+    def visit_mypy_file(self, mypyfile: MypyFile) -> None:
+        assert False, "use transform_mypy_file instead"
+
+    def visit_class_def(self, cdef: ClassDef) -> None:
+        transform_class_def(self.builder, cdef)
+
+    def visit_import(self, node: Import) -> None:
+        transform_import(self.builder, node)
+
+    def visit_import_from(self, node: ImportFrom) -> None:
+        transform_import_from(self.builder, node)
+
+    def visit_import_all(self, node: ImportAll) -> None:
+        transform_import_all(self.builder, node)
+
+    def visit_func_def(self, fdef: FuncDef) -> None:
+        transform_func_def(self.builder, fdef)
+
+    def visit_overloaded_func_def(self, o: OverloadedFuncDef) -> None:
+        transform_overloaded_func_def(self.builder, o)
+
+    def visit_decorator(self, dec: Decorator) -> None:
+        transform_decorator(self.builder, dec)
+
+    def visit_block(self, block: Block) -> None:
+        transform_block(self.builder, block)
+
+    # Statements
+
+    def visit_expression_stmt(self, stmt: ExpressionStmt) -> None:
+        transform_expression_stmt(self.builder, stmt)
+
+    def visit_return_stmt(self, stmt: ReturnStmt) -> None:
+        transform_return_stmt(self.builder, stmt)
+
+    def visit_assignment_stmt(self, stmt: AssignmentStmt) -> None:
+        transform_assignment_stmt(self.builder, stmt)
+
+    def visit_operator_assignment_stmt(self, stmt: OperatorAssignmentStmt) -> None:
+        transform_operator_assignment_stmt(self.builder, stmt)
+
+    def visit_if_stmt(self, stmt: IfStmt) -> None:
+        transform_if_stmt(self.builder, stmt)
+
+    def visit_while_stmt(self, stmt: WhileStmt) -> None:
+        transform_while_stmt(self.builder, stmt)
+
+    def visit_for_stmt(self, stmt: ForStmt) -> None:
+        transform_for_stmt(self.builder, stmt)
+
+    def visit_break_stmt(self, stmt: BreakStmt) -> None:
+        transform_break_stmt(self.builder, stmt)
+
+    def visit_continue_stmt(self, stmt: ContinueStmt) -> None:
+        transform_continue_stmt(self.builder, stmt)
+
+    def visit_raise_stmt(self, stmt: RaiseStmt) -> None:
+        transform_raise_stmt(self.builder, stmt)
+
+    def visit_try_stmt(self, stmt: TryStmt) -> None:
+        transform_try_stmt(self.builder, stmt)
+
+    def visit_with_stmt(self, stmt: WithStmt) -> None:
+        transform_with_stmt(self.builder, stmt)
+
+    def visit_pass_stmt(self, stmt: PassStmt) -> None:
+        pass
+
+    def visit_assert_stmt(self, stmt: AssertStmt) -> None:
+        transform_assert_stmt(self.builder, stmt)
+
+    def visit_del_stmt(self, stmt: DelStmt) -> None:
+        transform_del_stmt(self.builder, stmt)
+
+    def visit_global_decl(self, stmt: GlobalDecl) -> None:
+        # Pure declaration -- no runtime effect
+        pass
+
+    def visit_nonlocal_decl(self, stmt: NonlocalDecl) -> None:
+        # Pure declaration -- no runtime effect
+        pass
+
+    def visit_match_stmt(self, stmt: MatchStmt) -> None:
+        transform_match_stmt(self.builder, stmt)
+
+    # Expressions
+
+    def visit_name_expr(self, expr: NameExpr) -> Value:
+        return transform_name_expr(self.builder, expr)
+
+    def visit_member_expr(self, expr: MemberExpr) -> Value:
+        return transform_member_expr(self.builder, expr)
+
+    def visit_super_expr(self, expr: SuperExpr) -> Value:
+        return transform_super_expr(self.builder, expr)
+
+    def visit_call_expr(self, expr: CallExpr) -> Value:
+        return transform_call_expr(self.builder, expr)
+
+    def visit_unary_expr(self, expr: UnaryExpr) -> Value:
+        return transform_unary_expr(self.builder, expr)
+
+    def visit_op_expr(self, expr: OpExpr) -> Value:
+        return transform_op_expr(self.builder, expr)
+
+    def visit_index_expr(self, expr: IndexExpr) -> Value:
+        return transform_index_expr(self.builder, expr)
+
+    def visit_conditional_expr(self, expr: ConditionalExpr) -> Value:
+        return transform_conditional_expr(self.builder, expr)
+
+    def visit_comparison_expr(self, expr: ComparisonExpr) -> Value:
+        return transform_comparison_expr(self.builder, expr)
+
+    def visit_int_expr(self, expr: IntExpr) -> Value:
+        return transform_int_expr(self.builder, expr)
+
+    def visit_float_expr(self, expr: FloatExpr) -> Value:
+        return transform_float_expr(self.builder, expr)
+
+    def visit_complex_expr(self, expr: ComplexExpr) -> Value:
+        return transform_complex_expr(self.builder, expr)
+
+    def visit_str_expr(self, expr: StrExpr) -> Value:
+        return transform_str_expr(self.builder, expr)
+
+    def visit_bytes_expr(self, expr: BytesExpr) -> Value:
+        return transform_bytes_expr(self.builder, expr)
+
+    def visit_ellipsis(self, expr: EllipsisExpr) -> Value:
+        return transform_ellipsis(self.builder, expr)
+
+    def visit_list_expr(self, expr: ListExpr) -> Value:
+        return transform_list_expr(self.builder, expr)
+
+    def visit_tuple_expr(self, expr: TupleExpr) -> Value:
+        return transform_tuple_expr(self.builder, expr)
+
+    def visit_dict_expr(self, expr: DictExpr) -> Value:
+        return transform_dict_expr(self.builder, expr)
+
+    def visit_set_expr(self, expr: SetExpr) -> Value:
+        return transform_set_expr(self.builder, expr)
+
+    def visit_list_comprehension(self, expr: ListComprehension) -> Value:
+        return transform_list_comprehension(self.builder, expr)
+
+    def visit_set_comprehension(self, expr: SetComprehension) -> Value:
+        return transform_set_comprehension(self.builder, expr)
+
+    def visit_dictionary_comprehension(self, expr: DictionaryComprehension) -> Value:
+        return transform_dictionary_comprehension(self.builder, expr)
+
+    def visit_slice_expr(self, expr: SliceExpr) -> Value:
+        return transform_slice_expr(self.builder, expr)
+
+    def visit_generator_expr(self, expr: GeneratorExpr) -> Value:
+        return transform_generator_expr(self.builder, expr)
+
+    def visit_lambda_expr(self, expr: LambdaExpr) -> Value:
+        return transform_lambda_expr(self.builder, expr)
+
+    def visit_yield_expr(self, expr: YieldExpr) -> Value:
+        return transform_yield_expr(self.builder, expr)
+
+    def visit_yield_from_expr(self, o: YieldFromExpr) -> Value:
+        return transform_yield_from_expr(self.builder, o)
+
+    def visit_await_expr(self, o: AwaitExpr) -> Value:
+        return transform_await_expr(self.builder, o)
+
+    def visit_assignment_expr(self, o: AssignmentExpr) -> Value:
+        return transform_assignment_expr(self.builder, o)
+
+    # Constructs that shouldn't ever show up
+
+    def visit_enum_call_expr(self, o: EnumCallExpr) -> Value:
+        assert False, "can't compile analysis-only expressions"
+
+    def visit__promote_expr(self, o: PromoteExpr) -> Value:
+        assert False, "can't compile analysis-only expressions"
+
+    def visit_namedtuple_expr(self, o: NamedTupleExpr) -> Value:
+        assert False, "can't compile analysis-only expressions"
+
+    def visit_newtype_expr(self, o: NewTypeExpr) -> Value:
+        assert False, "can't compile analysis-only expressions"
+
+    def visit_temp_node(self, o: TempNode) -> Value:
+        assert False, "can't compile analysis-only expressions"
+
+    def visit_type_alias_expr(self, o: TypeAliasExpr) -> Value:
+        assert False, "can't compile analysis-only expressions"
+
+    def visit_type_application(self, o: TypeApplication) -> Value:
+        assert False, "can't compile analysis-only expressions"
+
+    def visit_type_var_expr(self, o: TypeVarExpr) -> Value:
+        assert False, "can't compile analysis-only expressions"
+
+    def visit_paramspec_expr(self, o: ParamSpecExpr) -> Value:
+        assert False, "can't compile analysis-only expressions"
+
+    def visit_type_var_tuple_expr(self, o: TypeVarTupleExpr) -> Value:
+        assert False, "can't compile analysis-only expressions"
+
+    def visit_typeddict_expr(self, o: TypedDictExpr) -> Value:
+        assert False, "can't compile analysis-only expressions"
+
+    def visit_reveal_expr(self, o: RevealExpr) -> Value:
+        assert False, "can't compile analysis-only expressions"
+
+    def visit_var(self, o: Var) -> None:
+        assert False, "can't compile Var; should have been handled already?"
+
+    def visit_cast_expr(self, o: CastExpr) -> Value:
+        assert False, "CastExpr should have been handled in CallExpr"
+
+    def visit_assert_type_expr(self, o: AssertTypeExpr) -> Value:
+        assert False, "AssertTypeExpr should have been handled in CallExpr"
+
+    def visit_star_expr(self, o: StarExpr) -> Value:
+        assert False, "should have been handled in Tuple/List/Set/DictExpr or CallExpr"
+
+    # Helpers
+
+    def bail(self, msg: str, line: int) -> NoReturn:
+        """Reports an error and aborts compilation up until the last accept() call
+
+        (accept() catches the UnsupportedException and keeps on
+        processing. This allows errors to be non-blocking without always
+        needing to write handling for them.
+        """
+        self.builder.error(msg, line)
+        raise UnsupportedException()

venv/lib/python3.12/site-packages/mypyc/irbuild/vtable.py

@@ -0,0 +1,82 @@
+"""Compute vtables of native (extension) classes."""
+
+from __future__ import annotations
+
+import itertools
+
+from mypyc.ir.class_ir import ClassIR, VTableEntries, VTableMethod
+from mypyc.sametype import is_same_method_signature
+
+
+def compute_vtable(cls: ClassIR) -> None:
+    """Compute the vtable structure for a class."""
+    if cls.vtable is not None:
+        return
+
+    if not cls.is_generated:
+        cls.has_dict = any(x.inherits_python for x in cls.mro)
+
+    for t in cls.mro[1:]:
+        # Make sure all ancestors are processed first
+        compute_vtable(t)
+        # Merge attributes from traits into the class
+        if not t.is_trait:
+            continue
+        for name, typ in t.attributes.items():
+            if not cls.is_trait and not any(name in b.attributes for b in cls.base_mro):
+                cls.attributes[name] = typ
+
+    cls.vtable = {}
+    if cls.base:
+        assert cls.base.vtable is not None
+        cls.vtable.update(cls.base.vtable)
+        cls.vtable_entries = specialize_parent_vtable(cls, cls.base)
+
+    # Include the vtable from the parent classes, but handle method overrides.
+    entries = cls.vtable_entries
+
+    all_traits = [t for t in cls.mro if t.is_trait]
+
+    for t in [cls] + cls.traits:
+        for fn in itertools.chain(t.methods.values()):
+            # TODO: don't generate a new entry when we overload without changing the type
+            if fn == cls.get_method(fn.name, prefer_method=True):
+                cls.vtable[fn.name] = len(entries)
+                # If the class contains a glue method referring to itself, that is a
+                # shadow glue method to support interpreted subclasses.
+                shadow = cls.glue_methods.get((cls, fn.name))
+                entries.append(VTableMethod(t, fn.name, fn, shadow))
+
+    # Compute vtables for all of the traits that the class implements
+    if not cls.is_trait:
+        for trait in all_traits:
+            compute_vtable(trait)
+            cls.trait_vtables[trait] = specialize_parent_vtable(cls, trait)
+
+
+def specialize_parent_vtable(cls: ClassIR, parent: ClassIR) -> VTableEntries:
+    """Generate the part of a vtable corresponding to a parent class or trait"""
+    updated = []
+    for entry in parent.vtable_entries:
+        # Find the original method corresponding to this vtable entry.
+        # (This may not be the method in the entry, if it was overridden.)
+        orig_parent_method = entry.cls.get_method(entry.name, prefer_method=True)
+        assert orig_parent_method
+        method_cls = cls.get_method_and_class(entry.name, prefer_method=True)
+        if method_cls:
+            child_method, defining_cls = method_cls
+            # TODO: emit a wrapper for __init__ that raises or something
+            if (
+                is_same_method_signature(orig_parent_method.sig, child_method.sig)
+                or orig_parent_method.name == "__init__"
+            ):
+                entry = VTableMethod(entry.cls, entry.name, child_method, entry.shadow_method)
+            else:
+                entry = VTableMethod(
+                    entry.cls,
+                    entry.name,
+                    defining_cls.glue_methods[(entry.cls, entry.name)],
+                    entry.shadow_method,
+                )
+        updated.append(entry)
+    return updated