[mypyc] Implement async with (#13442)

Also fix returning a value from inside a try block when the finally block
does a yield. (Which happens when returning from an async with).

Progress on mypyc/mypyc##868.

mypyc/irbuild/builder.py

@@ -786,11 +786,16 @@ def push_loop_stack(self, continue_block: BasicBlock, break_block: BasicBlock) -
     def pop_loop_stack(self) -> None:
         self.nonlocal_control.pop()
 
-    def spill(self, value: Value) -> AssignmentTarget:
+    def make_spill_target(self, type: RType) -> AssignmentTarget:
         """Moves a given Value instance into the generator class' environment class."""
         name = f"{TEMP_ATTR_NAME}{self.temp_counter}"
         self.temp_counter += 1
-        target = self.add_var_to_env_class(Var(name), value.type, self.fn_info.generator_class)
+        target = self.add_var_to_env_class(Var(name), type, self.fn_info.generator_class)
+        return target
+
+    def spill(self, value: Value) -> AssignmentTarget:
+        """Moves a given Value instance into the generator class' environment class."""
+        target = self.make_spill_target(value.type)
         # Shouldn't be able to fail, so -1 for line
         self.assign(target, value, -1)
         return target

mypyc/irbuild/function.py

@@ -17,7 +17,6 @@
 
 from mypy.nodes import (
     ArgKind,
-    AwaitExpr,
     ClassDef,
     Decorator,
     FuncDef,
@@ -27,8 +26,6 @@
     SymbolNode,
     TypeInfo,
     Var,
-    YieldExpr,
-    YieldFromExpr,
 )
 from mypy.types import CallableType, get_proper_type
 from mypyc.common import LAMBDA_NAME, SELF_NAME
@@ -44,7 +41,6 @@
 )
 from mypyc.ir.ops import (
     BasicBlock,
-    Branch,
     GetAttr,
     InitStatic,
     Integer,
@@ -62,7 +58,6 @@
     bool_rprimitive,
     dict_rprimitive,
     int_rprimitive,
-    object_pointer_rprimitive,
     object_rprimitive,
 )
 from mypyc.irbuild.builder import IRBuilder, SymbolTarget, gen_arg_defaults
@@ -88,18 +83,11 @@
     populate_switch_for_generator_class,
     setup_env_for_generator_class,
 )
-from mypyc.irbuild.statement import transform_try_except
 from mypyc.irbuild.targets import AssignmentTarget
 from mypyc.irbuild.util import is_constant
 from mypyc.primitives.dict_ops import dict_get_method_with_none, dict_new_op, dict_set_item_op
-from mypyc.primitives.generic_ops import iter_op, next_raw_op, py_setattr_op
-from mypyc.primitives.misc_ops import (
-    check_stop_op,
-    coro_op,
-    register_function,
-    send_op,
-    yield_from_except_op,
-)
+from mypyc.primitives.generic_ops import py_setattr_op
+from mypyc.primitives.misc_ops import register_function
 from mypyc.primitives.registry import builtin_names
 from mypyc.sametype import is_same_method_signature
 
@@ -178,25 +166,6 @@ def transform_lambda_expr(builder: IRBuilder, expr: LambdaExpr) -> Value:
     return func_reg
 
 
-def transform_yield_expr(builder: IRBuilder, expr: YieldExpr) -> Value:
-    if builder.fn_info.is_coroutine:
-        builder.error("async generators are unimplemented", expr.line)
-
-    if expr.expr:
-        retval = builder.accept(expr.expr)
-    else:
-        retval = builder.builder.none()
-    return emit_yield(builder, retval, expr.line)
-
-
-def transform_yield_from_expr(builder: IRBuilder, o: YieldFromExpr) -> Value:
-    return handle_yield_from_and_await(builder, o)
-
-
-def transform_await_expr(builder: IRBuilder, o: AwaitExpr) -> Value:
-    return handle_yield_from_and_await(builder, o)
-
-
 # Internal functions
 
 
@@ -551,112 +520,6 @@ def gen_func_ns(builder: IRBuilder) -> str:
     )
 
 
-def emit_yield(builder: IRBuilder, val: Value, line: int) -> Value:
-    retval = builder.coerce(val, builder.ret_types[-1], line)
-
-    cls = builder.fn_info.generator_class
-    # Create a new block for the instructions immediately following the yield expression, and
-    # set the next label so that the next time '__next__' is called on the generator object,
-    # the function continues at the new block.
-    next_block = BasicBlock()
-    next_label = len(cls.continuation_blocks)
-    cls.continuation_blocks.append(next_block)
-    builder.assign(cls.next_label_target, Integer(next_label), line)
-    builder.add(Return(retval))
-    builder.activate_block(next_block)
-
-    add_raise_exception_blocks_to_generator_class(builder, line)
-
-    assert cls.send_arg_reg is not None
-    return cls.send_arg_reg
-
-
-def handle_yield_from_and_await(builder: IRBuilder, o: Union[YieldFromExpr, AwaitExpr]) -> Value:
-    # This is basically an implementation of the code in PEP 380.
-
-    # TODO: do we want to use the right types here?
-    result = Register(object_rprimitive)
-    to_yield_reg = Register(object_rprimitive)
-    received_reg = Register(object_rprimitive)
-
-    if isinstance(o, YieldFromExpr):
-        iter_val = builder.call_c(iter_op, [builder.accept(o.expr)], o.line)
-    else:
-        iter_val = builder.call_c(coro_op, [builder.accept(o.expr)], o.line)
-
-    iter_reg = builder.maybe_spill_assignable(iter_val)
-
-    stop_block, main_block, done_block = BasicBlock(), BasicBlock(), BasicBlock()
-    _y_init = builder.call_c(next_raw_op, [builder.read(iter_reg)], o.line)
-    builder.add(Branch(_y_init, stop_block, main_block, Branch.IS_ERROR))
-
-    # Try extracting a return value from a StopIteration and return it.
-    # If it wasn't, this reraises the exception.
-    builder.activate_block(stop_block)
-    builder.assign(result, builder.call_c(check_stop_op, [], o.line), o.line)
-    builder.goto(done_block)
-
-    builder.activate_block(main_block)
-    builder.assign(to_yield_reg, _y_init, o.line)
-
-    # OK Now the main loop!
-    loop_block = BasicBlock()
-    builder.goto_and_activate(loop_block)
-
-    def try_body() -> None:
-        builder.assign(
-            received_reg, emit_yield(builder, builder.read(to_yield_reg), o.line), o.line
-        )
-
-    def except_body() -> None:
-        # The body of the except is all implemented in a C function to
-        # reduce how much code we need to generate. It returns a value
-        # indicating whether to break or yield (or raise an exception).
-        val = Register(object_rprimitive)
-        val_address = builder.add(LoadAddress(object_pointer_rprimitive, val))
-        to_stop = builder.call_c(
-            yield_from_except_op, [builder.read(iter_reg), val_address], o.line
-        )
-
-        ok, stop = BasicBlock(), BasicBlock()
-        builder.add(Branch(to_stop, stop, ok, Branch.BOOL))
-
-        # The exception got swallowed. Continue, yielding the returned value
-        builder.activate_block(ok)
-        builder.assign(to_yield_reg, val, o.line)
-        builder.nonlocal_control[-1].gen_continue(builder, o.line)
-
-        # The exception was a StopIteration. Stop iterating.
-        builder.activate_block(stop)
-        builder.assign(result, val, o.line)
-        builder.nonlocal_control[-1].gen_break(builder, o.line)
-
-    def else_body() -> None:
-        # Do a next() or a .send(). It will return NULL on exception
-        # but it won't automatically propagate.
-        _y = builder.call_c(send_op, [builder.read(iter_reg), builder.read(received_reg)], o.line)
-        ok, stop = BasicBlock(), BasicBlock()
-        builder.add(Branch(_y, stop, ok, Branch.IS_ERROR))
-
-        # Everything's fine. Yield it.
-        builder.activate_block(ok)
-        builder.assign(to_yield_reg, _y, o.line)
-        builder.nonlocal_control[-1].gen_continue(builder, o.line)
-
-        # Try extracting a return value from a StopIteration and return it.
-        # If it wasn't, this rereaises the exception.
-        builder.activate_block(stop)
-        builder.assign(result, builder.call_c(check_stop_op, [], o.line), o.line)
-        builder.nonlocal_control[-1].gen_break(builder, o.line)
-
-    builder.push_loop_stack(loop_block, done_block)
-    transform_try_except(builder, try_body, [(None, None, except_body)], else_body, o.line)
-    builder.pop_loop_stack()
-
-    builder.goto_and_activate(done_block)
-    return builder.read(result)
-
-
 def load_decorated_func(builder: IRBuilder, fdef: FuncDef, orig_func_reg: Value) -> Value:
     """Apply decorators to a function.
 

mypyc/irbuild/nonlocalcontrol.py

@@ -6,11 +6,10 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Union
 
 from mypyc.ir.ops import (
     NO_TRACEBACK_LINE_NO,
-    Assign,
     BasicBlock,
     Branch,
     Goto,
@@ -142,7 +141,7 @@ class TryFinallyNonlocalControl(NonlocalControl):
 
     def __init__(self, target: BasicBlock) -> None:
         self.target = target
-        self.ret_reg: Optional[Register] = None
+        self.ret_reg: Union[None, Register, AssignmentTarget] = None
 
     def gen_break(self, builder: IRBuilder, line: int) -> None:
         builder.error("break inside try/finally block is unimplemented", line)
@@ -152,9 +151,15 @@ def gen_continue(self, builder: IRBuilder, line: int) -> None:
 
     def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
         if self.ret_reg is None:
-            self.ret_reg = Register(builder.ret_types[-1])
+            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(Assign(self.ret_reg, value))
         builder.add(Goto(self.target))
 
 
@@ -180,9 +185,8 @@ class FinallyNonlocalControl(CleanupNonlocalControl):
     leave and the return register is decrefed if it isn't null.
     """
 
-    def __init__(self, outer: NonlocalControl, ret_reg: Optional[Value], saved: Value) -> None:
+    def __init__(self, outer: NonlocalControl, saved: Value) -> None:
         super().__init__(outer)
-        self.ret_reg = ret_reg
         self.saved = saved
 
     def gen_cleanup(self, builder: IRBuilder, line: int) -> None: