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

mypyc/irbuild/match.py

@@ -0,0 +1,355 @@
+from contextlib import contextmanager
+from typing import Generator, List, Optional, Tuple
+
+from mypy.nodes import MatchStmt, NameExpr, TypeInfo
+from mypy.patterns import (
+    AsPattern,
+    ClassPattern,
+    MappingPattern,
+    OrPattern,
+    Pattern,
+    SequencePattern,
+    SingletonPattern,
+    StarredPattern,
+    ValuePattern,
+)
+from mypy.traverser import TraverserVisitor
+from mypy.types import Instance, TupleType, get_proper_type
+from mypyc.ir.ops import BasicBlock, Value
+from mypyc.ir.rtypes import object_rprimitive
+from mypyc.irbuild.builder import IRBuilder
+from mypyc.primitives.dict_ops import (
+    dict_copy,
+    dict_del_item,
+    mapping_has_key,
+    supports_mapping_protocol,
+)
+from mypyc.primitives.generic_ops import generic_ssize_t_len_op
+from mypyc.primitives.list_ops import (
+    sequence_get_item,
+    sequence_get_slice,
+    supports_sequence_protocol,
+)
+from mypyc.primitives.misc_ops import fast_isinstance_op, slow_isinstance_op
+
+# From: https://peps.python.org/pep-0634/#class-patterns
+MATCHABLE_BUILTINS = {
+    "builtins.bool",
+    "builtins.bytearray",
+    "builtins.bytes",
+    "builtins.dict",
+    "builtins.float",
+    "builtins.frozenset",
+    "builtins.int",
+    "builtins.list",
+    "builtins.set",
+    "builtins.str",
+    "builtins.tuple",
+}
+
+
+class MatchVisitor(TraverserVisitor):
+    builder: IRBuilder
+    code_block: BasicBlock
+    next_block: BasicBlock
+    final_block: BasicBlock
+    subject: Value
+    match: MatchStmt
+
+    as_pattern: Optional[AsPattern] = None
+
+    def __init__(self, builder: IRBuilder, match_node: MatchStmt) -> None:
+        self.builder = builder
+
+        self.code_block = BasicBlock()
+        self.next_block = BasicBlock()
+        self.final_block = BasicBlock()
+
+        self.match = match_node
+        self.subject = builder.accept(match_node.subject)
+
+    def build_match_body(self, index: int) -> None:
+        self.builder.activate_block(self.code_block)
+
+        guard = self.match.guards[index]
+
+        if guard:
+            self.code_block = BasicBlock()
+
+            cond = self.builder.accept(guard)
+            self.builder.add_bool_branch(cond, self.code_block, self.next_block)
+
+            self.builder.activate_block(self.code_block)
+
+        self.builder.accept(self.match.bodies[index])
+        self.builder.goto(self.final_block)
+
+    def visit_match_stmt(self, m: MatchStmt) -> None:
+        for i, pattern in enumerate(m.patterns):
+            self.code_block = BasicBlock()
+            self.next_block = BasicBlock()
+
+            pattern.accept(self)
+
+            self.build_match_body(i)
+            self.builder.activate_block(self.next_block)
+
+        self.builder.goto_and_activate(self.final_block)
+
+    def visit_value_pattern(self, pattern: ValuePattern) -> None:
+        value = self.builder.accept(pattern.expr)
+
+        cond = self.builder.binary_op(self.subject, value, "==", pattern.expr.line)
+
+        self.bind_as_pattern(value)
+
+        self.builder.add_bool_branch(cond, self.code_block, self.next_block)
+
+    def visit_or_pattern(self, pattern: OrPattern) -> None:
+        backup_block = self.next_block
+        self.next_block = BasicBlock()
+
+        for p in pattern.patterns:
+            # Hack to ensure the as pattern is bound to each pattern in the
+            # "or" pattern, but not every subpattern
+            backup = self.as_pattern
+            p.accept(self)
+            self.as_pattern = backup
+
+            self.builder.activate_block(self.next_block)
+            self.next_block = BasicBlock()
+
+        self.next_block = backup_block
+        self.builder.goto(self.next_block)
+
+    def visit_class_pattern(self, pattern: ClassPattern) -> None:
+        # TODO: use faster instance check for native classes (while still
+        # making sure to account for inheritence)
+        isinstance_op = (
+            fast_isinstance_op
+            if self.builder.is_builtin_ref_expr(pattern.class_ref)
+            else slow_isinstance_op
+        )
+
+        cond = self.builder.call_c(
+            isinstance_op, [self.subject, self.builder.accept(pattern.class_ref)], pattern.line
+        )
+
+        self.builder.add_bool_branch(cond, self.code_block, self.next_block)
+
+        self.bind_as_pattern(self.subject, new_block=True)
+
+        if pattern.positionals:
+            if pattern.class_ref.fullname in MATCHABLE_BUILTINS:
+                self.builder.activate_block(self.code_block)
+                self.code_block = BasicBlock()
+
+                pattern.positionals[0].accept(self)
+
+                return
+
+            node = pattern.class_ref.node
+            assert isinstance(node, TypeInfo)
+
+            ty = node.names.get("__match_args__")
+            assert ty
+
+            match_args_type = get_proper_type(ty.type)
+            assert isinstance(match_args_type, TupleType)
+
+            match_args: List[str] = []
+
+            for item in match_args_type.items:
+                proper_item = get_proper_type(item)
+                assert isinstance(proper_item, Instance) and proper_item.last_known_value
+
+                match_arg = proper_item.last_known_value.value
+                assert isinstance(match_arg, str)
+
+                match_args.append(match_arg)
+
+            for i, expr in enumerate(pattern.positionals):
+                self.builder.activate_block(self.code_block)
+                self.code_block = BasicBlock()
+
+                # TODO: use faster "get_attr" method instead when calling on native or
+                # builtin objects
+                positional = self.builder.py_get_attr(self.subject, match_args[i], expr.line)
+
+                with self.enter_subpattern(positional):
+                    expr.accept(self)
+
+        for key, value in zip(pattern.keyword_keys, pattern.keyword_values):
+            self.builder.activate_block(self.code_block)
+            self.code_block = BasicBlock()
+
+            # TODO: same as above "get_attr" comment
+            attr = self.builder.py_get_attr(self.subject, key, value.line)
+
+            with self.enter_subpattern(attr):
+                value.accept(self)
+
+    def visit_as_pattern(self, pattern: AsPattern) -> None:
+        if pattern.pattern:
+            old_pattern = self.as_pattern
+            self.as_pattern = pattern
+            pattern.pattern.accept(self)
+            self.as_pattern = old_pattern
+
+        elif pattern.name:
+            target = self.builder.get_assignment_target(pattern.name)
+
+            self.builder.assign(target, self.subject, pattern.line)
+
+        self.builder.goto(self.code_block)
+
+    def visit_singleton_pattern(self, pattern: SingletonPattern) -> None:
+        if pattern.value is None:
+            obj = self.builder.none_object()
+        elif pattern.value is True:
+            obj = self.builder.true()
+        else:
+            obj = self.builder.false()
+
+        cond = self.builder.binary_op(self.subject, obj, "is", pattern.line)
+
+        self.builder.add_bool_branch(cond, self.code_block, self.next_block)
+
+    def visit_mapping_pattern(self, pattern: MappingPattern) -> None:
+        is_dict = self.builder.call_c(supports_mapping_protocol, [self.subject], pattern.line)
+
+        self.builder.add_bool_branch(is_dict, self.code_block, self.next_block)
+
+        keys: List[Value] = []
+
+        for key, value in zip(pattern.keys, pattern.values):
+            self.builder.activate_block(self.code_block)
+            self.code_block = BasicBlock()
+
+            key_value = self.builder.accept(key)
+            keys.append(key_value)
+
+            exists = self.builder.call_c(mapping_has_key, [self.subject, key_value], pattern.line)
+
+            self.builder.add_bool_branch(exists, self.code_block, self.next_block)
+            self.builder.activate_block(self.code_block)
+            self.code_block = BasicBlock()
+
+            item = self.builder.gen_method_call(
+                self.subject, "__getitem__", [key_value], object_rprimitive, pattern.line
+            )
+
+            with self.enter_subpattern(item):
+                value.accept(self)
+
+        if pattern.rest:
+            self.builder.activate_block(self.code_block)
+            self.code_block = BasicBlock()
+
+            rest = self.builder.call_c(dict_copy, [self.subject], pattern.rest.line)
+
+            target = self.builder.get_assignment_target(pattern.rest)
+
+            self.builder.assign(target, rest, pattern.rest.line)
+
+            for i, key_name in enumerate(keys):
+                self.builder.call_c(dict_del_item, [rest, key_name], pattern.keys[i].line)
+
+            self.builder.goto(self.code_block)
+
+    def visit_sequence_pattern(self, seq_pattern: SequencePattern) -> None:
+        star_index, capture, patterns = prep_sequence_pattern(seq_pattern)
+
+        is_list = self.builder.call_c(supports_sequence_protocol, [self.subject], seq_pattern.line)
+
+        self.builder.add_bool_branch(is_list, self.code_block, self.next_block)
+
+        self.builder.activate_block(self.code_block)
+        self.code_block = BasicBlock()
+
+        actual_len = self.builder.call_c(generic_ssize_t_len_op, [self.subject], seq_pattern.line)
+        min_len = len(patterns)
+
+        is_long_enough = self.builder.binary_op(
+            actual_len,
+            self.builder.load_int(min_len),
+            "==" if star_index is None else ">=",
+            seq_pattern.line,
+        )
+
+        self.builder.add_bool_branch(is_long_enough, self.code_block, self.next_block)
+
+        for i, pattern in enumerate(patterns):
+            self.builder.activate_block(self.code_block)
+            self.code_block = BasicBlock()
+
+            if star_index is not None and i >= star_index:
+                current = self.builder.binary_op(
+                    actual_len, self.builder.load_int(min_len - i), "-", pattern.line
+                )
+
+            else:
+                current = self.builder.load_int(i)
+
+            item = self.builder.call_c(sequence_get_item, [self.subject, current], pattern.line)
+
+            with self.enter_subpattern(item):
+                pattern.accept(self)
+
+        if capture and star_index is not None:
+            self.builder.activate_block(self.code_block)
+            self.code_block = BasicBlock()
+
+            capture_end = self.builder.binary_op(
+                actual_len, self.builder.load_int(min_len - star_index), "-", capture.line
+            )
+
+            rest = self.builder.call_c(
+                sequence_get_slice,
+                [self.subject, self.builder.load_int(star_index), capture_end],
+                capture.line,
+            )
+
+            target = self.builder.get_assignment_target(capture)
+            self.builder.assign(target, rest, capture.line)
+
+            self.builder.goto(self.code_block)
+
+    def bind_as_pattern(self, value: Value, new_block: bool = False) -> None:
+        if self.as_pattern and self.as_pattern.pattern and self.as_pattern.name:
+            if new_block:
+                self.builder.activate_block(self.code_block)
+                self.code_block = BasicBlock()
+
+            target = self.builder.get_assignment_target(self.as_pattern.name)
+            self.builder.assign(target, value, self.as_pattern.pattern.line)
+
+            self.as_pattern = None
+
+            if new_block:
+                self.builder.goto(self.code_block)
+
+    @contextmanager
+    def enter_subpattern(self, subject: Value) -> Generator[None, None, None]:
+        old_subject = self.subject
+        self.subject = subject
+        yield
+        self.subject = old_subject
+
+
+def prep_sequence_pattern(
+    seq_pattern: SequencePattern,
+) -> Tuple[Optional[int], Optional[NameExpr], List[Pattern]]:
+    star_index: Optional[int] = None
+    capture: Optional[NameExpr] = None
+    patterns: List[Pattern] = []
+
+    for i, pattern in enumerate(seq_pattern.patterns):
+        if isinstance(pattern, StarredPattern):
+            star_index = i
+            capture = pattern.capture
+
+        else:
+            patterns.append(pattern)
+
+    return star_index, capture, patterns

mypyc/irbuild/nonlocalcontrol.py

@@ -0,0 +1,197 @@
+"""Helpers for dealing with nonlocal control such as 'break' and 'return'.
+
+Model how these behave differently in different contexts.
+"""
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING
+
+from mypyc.ir.ops import (
+    NO_TRACEBACK_LINE_NO,
+    BasicBlock,
+    Branch,
+    Goto,
+    Integer,
+    Register,
+    Return,
+    Unreachable,
+    Value,
+)
+from mypyc.irbuild.targets import AssignmentTarget
+from mypyc.primitives.exc_ops import restore_exc_info_op, set_stop_iteration_value
+
+if TYPE_CHECKING:
+    from mypyc.irbuild.builder import IRBuilder
+
+
+class NonlocalControl:
+    """ABC representing a stack frame of constructs that modify nonlocal control flow.
+
+    The nonlocal control flow constructs are break, continue, and
+    return, and their behavior is modified by a number of other
+    constructs.  The most obvious is loop, which override where break
+    and continue jump to, but also `except` (which needs to clear
+    exc_info when left) and (eventually) finally blocks (which need to
+    ensure that the finally block is always executed when leaving the
+    try/except blocks).
+    """
+
+    @abstractmethod
+    def gen_break(self, builder: IRBuilder, line: int) -> None:
+        pass
+
+    @abstractmethod
+    def gen_continue(self, builder: IRBuilder, line: int) -> None:
+        pass
+
+    @abstractmethod
+    def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
+        pass
+
+
+class BaseNonlocalControl(NonlocalControl):
+    """Default nonlocal control outside any statements that affect it."""
+
+    def gen_break(self, builder: IRBuilder, line: int) -> None:
+        assert False, "break outside of loop"
+
+    def gen_continue(self, builder: IRBuilder, line: int) -> None:
+        assert False, "continue outside of loop"
+
+    def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
+        builder.add(Return(value))
+
+
+class LoopNonlocalControl(NonlocalControl):
+    """Nonlocal control within a loop."""
+
+    def __init__(
+        self, outer: NonlocalControl, continue_block: BasicBlock, break_block: BasicBlock
+    ) -> None:
+        self.outer = outer
+        self.continue_block = continue_block
+        self.break_block = break_block
+
+    def gen_break(self, builder: IRBuilder, line: int) -> None:
+        builder.add(Goto(self.break_block))
+
+    def gen_continue(self, builder: IRBuilder, line: int) -> None:
+        builder.add(Goto(self.continue_block))
+
+    def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
+        self.outer.gen_return(builder, value, line)
+
+
+class GeneratorNonlocalControl(BaseNonlocalControl):
+    """Default nonlocal control in a generator function outside statements."""
+
+    def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
+        # Assign an invalid next label number so that the next time
+        # __next__ is called, we jump to the case in which
+        # StopIteration is raised.
+        builder.assign(builder.fn_info.generator_class.next_label_target, Integer(-1), line)
+
+        # Raise a StopIteration containing a field for the value that
+        # should be returned. Before doing so, create a new block
+        # without an error handler set so that the implicitly thrown
+        # StopIteration isn't caught by except blocks inside of the
+        # generator function.
+        builder.builder.push_error_handler(None)
+        builder.goto_and_activate(BasicBlock())
+
+        # Skip creating a traceback frame when we raise here, because
+        # we don't care about the traceback frame and it is kind of
+        # expensive since raising StopIteration is an extremely common
+        # case.  Also we call a special internal function to set
+        # StopIteration instead of using RaiseStandardError because
+        # the obvious thing doesn't work if the value is a tuple
+        # (???).
+        builder.call_c(set_stop_iteration_value, [value], NO_TRACEBACK_LINE_NO)
+        builder.add(Unreachable())
+        builder.builder.pop_error_handler()
+
+
+class CleanupNonlocalControl(NonlocalControl):
+    """Abstract nonlocal control that runs some cleanup code."""
+
+    def __init__(self, outer: NonlocalControl) -> None:
+        self.outer = outer
+
+    @abstractmethod
+    def gen_cleanup(self, builder: IRBuilder, line: int) -> None: ...
+
+    def gen_break(self, builder: IRBuilder, line: int) -> None:
+        self.gen_cleanup(builder, line)
+        self.outer.gen_break(builder, line)
+
+    def gen_continue(self, builder: IRBuilder, line: int) -> None:
+        self.gen_cleanup(builder, line)
+        self.outer.gen_continue(builder, line)
+
+    def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
+        self.gen_cleanup(builder, line)
+        self.outer.gen_return(builder, value, line)
+
+
+class TryFinallyNonlocalControl(NonlocalControl):
+    """Nonlocal control within try/finally."""
+
+    def __init__(self, target: BasicBlock) -> None:
+        self.target = target
+        self.ret_reg: None | Register | AssignmentTarget = None
+
+    def gen_break(self, builder: IRBuilder, line: int) -> None:
+        builder.error("break inside try/finally block is unimplemented", line)
+
+    def gen_continue(self, builder: IRBuilder, line: int) -> None:
+        builder.error("continue inside try/finally block is unimplemented", line)
+
+    def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
+        if self.ret_reg is None:
+            if builder.fn_info.is_generator:
+                self.ret_reg = builder.make_spill_target(builder.ret_types[-1])
+            else:
+                self.ret_reg = Register(builder.ret_types[-1])
+        # assert needed because of apparent mypy bug... it loses track of the union
+        # and infers the type as object
+        assert isinstance(self.ret_reg, (Register, AssignmentTarget))
+        builder.assign(self.ret_reg, value, line)
+
+        builder.add(Goto(self.target))
+
+
+class ExceptNonlocalControl(CleanupNonlocalControl):
+    """Nonlocal control for except blocks.
+
+    Just makes sure that sys.exc_info always gets restored when we leave.
+    This is super annoying.
+    """
+
+    def __init__(self, outer: NonlocalControl, saved: Value | AssignmentTarget) -> None:
+        super().__init__(outer)
+        self.saved = saved
+
+    def gen_cleanup(self, builder: IRBuilder, line: int) -> None:
+        builder.call_c(restore_exc_info_op, [builder.read(self.saved)], line)
+
+
+class FinallyNonlocalControl(CleanupNonlocalControl):
+    """Nonlocal control for finally blocks.
+
+    Just makes sure that sys.exc_info always gets restored when we
+    leave and the return register is decrefed if it isn't null.
+    """
+
+    def __init__(self, outer: NonlocalControl, saved: Value) -> None:
+        super().__init__(outer)
+        self.saved = saved
+
+    def gen_cleanup(self, builder: IRBuilder, line: int) -> None:
+        # Restore the old exc_info
+        target, cleanup = BasicBlock(), BasicBlock()
+        builder.add(Branch(self.saved, target, cleanup, Branch.IS_ERROR))
+        builder.activate_block(cleanup)
+        builder.call_c(restore_exc_info_op, [self.saved], line)
+        builder.goto_and_activate(target)