138 lines
4.4 KiB
Python
138 lines
4.4 KiB
Python
|
"""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 mypy.backports import OrderedDict
|
||
|
from typing import List, Dict, Callable, Any, TypeVar, cast
|
||
|
|
||
|
from mypy.nodes import MypyFile, Expression, ClassDef
|
||
|
from mypy.types import Type
|
||
|
from mypy.state import state
|
||
|
from mypy.build import Graph
|
||
|
|
||
|
from mypyc.common import TOP_LEVEL_NAME
|
||
|
from mypyc.errors import Errors
|
||
|
from mypyc.options import CompilerOptions
|
||
|
from mypyc.ir.rtypes import none_rprimitive
|
||
|
from mypyc.ir.module_ir import ModuleIR, ModuleIRs
|
||
|
from mypyc.ir.func_ir import FuncIR, FuncDecl, FuncSignature
|
||
|
from mypyc.irbuild.prebuildvisitor import PreBuildVisitor
|
||
|
from mypyc.irbuild.vtable import compute_vtable
|
||
|
from mypyc.irbuild.prepare import build_type_map, find_singledispatch_register_impls
|
||
|
from mypyc.irbuild.builder import IRBuilder
|
||
|
from mypyc.irbuild.visitor import IRBuilderVisitor
|
||
|
from mypyc.irbuild.mapper import Mapper
|
||
|
from mypyc.analysis.attrdefined import analyze_always_defined_attrs
|
||
|
|
||
|
|
||
|
# 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 IR for a set of modules that have been type-checked by mypy."""
|
||
|
|
||
|
build_type_map(mapper, modules, graph, types, options, errors)
|
||
|
singledispatch_info = find_singledispatch_register_impls(modules, errors)
|
||
|
|
||
|
result: ModuleIRs = OrderedDict()
|
||
|
|
||
|
# 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('<top level>')
|
||
|
|
||
|
# 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)
|