| """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 |