| """Maintain a mapping from mypy concepts to IR/compiled concepts.""" |
| |
| from typing import Dict, Optional |
| |
| from mypy.nodes import FuncDef, TypeInfo, SymbolNode, ARG_STAR, ARG_STAR2 |
| from mypy.types import ( |
| Instance, Type, CallableType, LiteralType, TypedDictType, UnboundType, PartialType, |
| UninhabitedType, Overloaded, UnionType, TypeType, AnyType, NoneTyp, TupleType, TypeVarType, |
| get_proper_type |
| ) |
| |
| from mypyc.ir.rtypes import ( |
| RType, RUnion, RTuple, RInstance, object_rprimitive, dict_rprimitive, tuple_rprimitive, |
| none_rprimitive, int_rprimitive, float_rprimitive, str_rprimitive, bool_rprimitive, |
| list_rprimitive, set_rprimitive |
| ) |
| from mypyc.ir.func_ir import FuncSignature, FuncDecl, RuntimeArg |
| from mypyc.ir.class_ir import ClassIR |
| |
| |
| 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, Optional[str]]) -> None: |
| self.group_map = group_map |
| self.type_to_ir = {} # type: Dict[TypeInfo, ClassIR] |
| self.func_to_decl = {} # type: Dict[SymbolNode, FuncDecl] |
| |
| def type_to_rtype(self, typ: Optional[Type]) -> 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.str': |
| return str_rprimitive |
| elif typ.type.fullname == 'builtins.bool': |
| return bool_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 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 |
| 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: int) -> 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)] |
| ret = self.type_to_rtype(fdef.type.ret_type) |
| else: |
| # Handle unannotated functions |
| arg_types = [object_rprimitive 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 |
| |
| args = [RuntimeArg(arg_name, arg_type, arg_kind) |
| for arg_name, arg_kind, arg_type in zip(fdef.arg_names, fdef.arg_kinds, arg_types)] |
| |
| # 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) |
