mypy-dev 1.9.0.dev0__py3-none-any.whl

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)

mypyc/irbuild/mapper.py

@@ -0,0 +1,221 @@
+"""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,
+    find_unpack_in_list,
+    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 or for variadic tuples.
+            if (
+                typ.partial_fallback.type.fullname == "builtins.tuple"
+                and find_unpack_in_list(typ.items) is None
+            ):
+                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