| """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, |
| int32_rprimitive, |
| int64_rprimitive, |
| int_rprimitive, |
| list_rprimitive, |
| none_rprimitive, |
| object_rprimitive, |
| range_rprimitive, |
| set_rprimitive, |
| str_rprimitive, |
| tuple_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 |
| 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([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 |
