Skip to content

Commit 3ec1849

Browse files
authored
[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.
1 parent 487b736 commit 3ec1849

11 files changed

+314
-231
lines changed

mypyc/irbuild/builder.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -786,11 +786,16 @@ def push_loop_stack(self, continue_block: BasicBlock, break_block: BasicBlock) -
786786
def pop_loop_stack(self) -> None:
787787
self.nonlocal_control.pop()
788788

789-
def spill(self, value: Value) -> AssignmentTarget:
789+
def make_spill_target(self, type: RType) -> AssignmentTarget:
790790
"""Moves a given Value instance into the generator class' environment class."""
791791
name = f"{TEMP_ATTR_NAME}{self.temp_counter}"
792792
self.temp_counter += 1
793-
target = self.add_var_to_env_class(Var(name), value.type, self.fn_info.generator_class)
793+
target = self.add_var_to_env_class(Var(name), type, self.fn_info.generator_class)
794+
return target
795+
796+
def spill(self, value: Value) -> AssignmentTarget:
797+
"""Moves a given Value instance into the generator class' environment class."""
798+
target = self.make_spill_target(value.type)
794799
# Shouldn't be able to fail, so -1 for line
795800
self.assign(target, value, -1)
796801
return target

mypyc/irbuild/function.py

Lines changed: 2 additions & 139 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717

1818
from mypy.nodes import (
1919
ArgKind,
20-
AwaitExpr,
2120
ClassDef,
2221
Decorator,
2322
FuncDef,
@@ -27,8 +26,6 @@
2726
SymbolNode,
2827
TypeInfo,
2928
Var,
30-
YieldExpr,
31-
YieldFromExpr,
3229
)
3330
from mypy.types import CallableType, get_proper_type
3431
from mypyc.common import LAMBDA_NAME, SELF_NAME
@@ -44,7 +41,6 @@
4441
)
4542
from mypyc.ir.ops import (
4643
BasicBlock,
47-
Branch,
4844
GetAttr,
4945
InitStatic,
5046
Integer,
@@ -62,7 +58,6 @@
6258
bool_rprimitive,
6359
dict_rprimitive,
6460
int_rprimitive,
65-
object_pointer_rprimitive,
6661
object_rprimitive,
6762
)
6863
from mypyc.irbuild.builder import IRBuilder, SymbolTarget, gen_arg_defaults
@@ -88,18 +83,11 @@
8883
populate_switch_for_generator_class,
8984
setup_env_for_generator_class,
9085
)
91-
from mypyc.irbuild.statement import transform_try_except
9286
from mypyc.irbuild.targets import AssignmentTarget
9387
from mypyc.irbuild.util import is_constant
9488
from mypyc.primitives.dict_ops import dict_get_method_with_none, dict_new_op, dict_set_item_op
95-
from mypyc.primitives.generic_ops import iter_op, next_raw_op, py_setattr_op
96-
from mypyc.primitives.misc_ops import (
97-
check_stop_op,
98-
coro_op,
99-
register_function,
100-
send_op,
101-
yield_from_except_op,
102-
)
89+
from mypyc.primitives.generic_ops import py_setattr_op
90+
from mypyc.primitives.misc_ops import register_function
10391
from mypyc.primitives.registry import builtin_names
10492
from mypyc.sametype import is_same_method_signature
10593

@@ -178,25 +166,6 @@ def transform_lambda_expr(builder: IRBuilder, expr: LambdaExpr) -> Value:
178166
return func_reg
179167

180168

181-
def transform_yield_expr(builder: IRBuilder, expr: YieldExpr) -> Value:
182-
if builder.fn_info.is_coroutine:
183-
builder.error("async generators are unimplemented", expr.line)
184-
185-
if expr.expr:
186-
retval = builder.accept(expr.expr)
187-
else:
188-
retval = builder.builder.none()
189-
return emit_yield(builder, retval, expr.line)
190-
191-
192-
def transform_yield_from_expr(builder: IRBuilder, o: YieldFromExpr) -> Value:
193-
return handle_yield_from_and_await(builder, o)
194-
195-
196-
def transform_await_expr(builder: IRBuilder, o: AwaitExpr) -> Value:
197-
return handle_yield_from_and_await(builder, o)
198-
199-
200169
# Internal functions
201170

202171

@@ -551,112 +520,6 @@ def gen_func_ns(builder: IRBuilder) -> str:
551520
)
552521

553522

554-
def emit_yield(builder: IRBuilder, val: Value, line: int) -> Value:
555-
retval = builder.coerce(val, builder.ret_types[-1], line)
556-
557-
cls = builder.fn_info.generator_class
558-
# Create a new block for the instructions immediately following the yield expression, and
559-
# set the next label so that the next time '__next__' is called on the generator object,
560-
# the function continues at the new block.
561-
next_block = BasicBlock()
562-
next_label = len(cls.continuation_blocks)
563-
cls.continuation_blocks.append(next_block)
564-
builder.assign(cls.next_label_target, Integer(next_label), line)
565-
builder.add(Return(retval))
566-
builder.activate_block(next_block)
567-
568-
add_raise_exception_blocks_to_generator_class(builder, line)
569-
570-
assert cls.send_arg_reg is not None
571-
return cls.send_arg_reg
572-
573-
574-
def handle_yield_from_and_await(builder: IRBuilder, o: Union[YieldFromExpr, AwaitExpr]) -> Value:
575-
# This is basically an implementation of the code in PEP 380.
576-
577-
# TODO: do we want to use the right types here?
578-
result = Register(object_rprimitive)
579-
to_yield_reg = Register(object_rprimitive)
580-
received_reg = Register(object_rprimitive)
581-
582-
if isinstance(o, YieldFromExpr):
583-
iter_val = builder.call_c(iter_op, [builder.accept(o.expr)], o.line)
584-
else:
585-
iter_val = builder.call_c(coro_op, [builder.accept(o.expr)], o.line)
586-
587-
iter_reg = builder.maybe_spill_assignable(iter_val)
588-
589-
stop_block, main_block, done_block = BasicBlock(), BasicBlock(), BasicBlock()
590-
_y_init = builder.call_c(next_raw_op, [builder.read(iter_reg)], o.line)
591-
builder.add(Branch(_y_init, stop_block, main_block, Branch.IS_ERROR))
592-
593-
# Try extracting a return value from a StopIteration and return it.
594-
# If it wasn't, this reraises the exception.
595-
builder.activate_block(stop_block)
596-
builder.assign(result, builder.call_c(check_stop_op, [], o.line), o.line)
597-
builder.goto(done_block)
598-
599-
builder.activate_block(main_block)
600-
builder.assign(to_yield_reg, _y_init, o.line)
601-
602-
# OK Now the main loop!
603-
loop_block = BasicBlock()
604-
builder.goto_and_activate(loop_block)
605-
606-
def try_body() -> None:
607-
builder.assign(
608-
received_reg, emit_yield(builder, builder.read(to_yield_reg), o.line), o.line
609-
)
610-
611-
def except_body() -> None:
612-
# The body of the except is all implemented in a C function to
613-
# reduce how much code we need to generate. It returns a value
614-
# indicating whether to break or yield (or raise an exception).
615-
val = Register(object_rprimitive)
616-
val_address = builder.add(LoadAddress(object_pointer_rprimitive, val))
617-
to_stop = builder.call_c(
618-
yield_from_except_op, [builder.read(iter_reg), val_address], o.line
619-
)
620-
621-
ok, stop = BasicBlock(), BasicBlock()
622-
builder.add(Branch(to_stop, stop, ok, Branch.BOOL))
623-
624-
# The exception got swallowed. Continue, yielding the returned value
625-
builder.activate_block(ok)
626-
builder.assign(to_yield_reg, val, o.line)
627-
builder.nonlocal_control[-1].gen_continue(builder, o.line)
628-
629-
# The exception was a StopIteration. Stop iterating.
630-
builder.activate_block(stop)
631-
builder.assign(result, val, o.line)
632-
builder.nonlocal_control[-1].gen_break(builder, o.line)
633-
634-
def else_body() -> None:
635-
# Do a next() or a .send(). It will return NULL on exception
636-
# but it won't automatically propagate.
637-
_y = builder.call_c(send_op, [builder.read(iter_reg), builder.read(received_reg)], o.line)
638-
ok, stop = BasicBlock(), BasicBlock()
639-
builder.add(Branch(_y, stop, ok, Branch.IS_ERROR))
640-
641-
# Everything's fine. Yield it.
642-
builder.activate_block(ok)
643-
builder.assign(to_yield_reg, _y, o.line)
644-
builder.nonlocal_control[-1].gen_continue(builder, o.line)
645-
646-
# Try extracting a return value from a StopIteration and return it.
647-
# If it wasn't, this rereaises the exception.
648-
builder.activate_block(stop)
649-
builder.assign(result, builder.call_c(check_stop_op, [], o.line), o.line)
650-
builder.nonlocal_control[-1].gen_break(builder, o.line)
651-
652-
builder.push_loop_stack(loop_block, done_block)
653-
transform_try_except(builder, try_body, [(None, None, except_body)], else_body, o.line)
654-
builder.pop_loop_stack()
655-
656-
builder.goto_and_activate(done_block)
657-
return builder.read(result)
658-
659-
660523
def load_decorated_func(builder: IRBuilder, fdef: FuncDef, orig_func_reg: Value) -> Value:
661524
"""Apply decorators to a function.
662525

mypyc/irbuild/nonlocalcontrol.py

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,10 @@
66
from __future__ import annotations
77

88
from abc import abstractmethod
9-
from typing import TYPE_CHECKING, Optional, Union
9+
from typing import TYPE_CHECKING, Union
1010

1111
from mypyc.ir.ops import (
1212
NO_TRACEBACK_LINE_NO,
13-
Assign,
1413
BasicBlock,
1514
Branch,
1615
Goto,
@@ -142,7 +141,7 @@ class TryFinallyNonlocalControl(NonlocalControl):
142141

143142
def __init__(self, target: BasicBlock) -> None:
144143
self.target = target
145-
self.ret_reg: Optional[Register] = None
144+
self.ret_reg: Union[None, Register, AssignmentTarget] = None
146145

147146
def gen_break(self, builder: IRBuilder, line: int) -> None:
148147
builder.error("break inside try/finally block is unimplemented", line)
@@ -152,9 +151,15 @@ def gen_continue(self, builder: IRBuilder, line: int) -> None:
152151

153152
def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
154153
if self.ret_reg is None:
155-
self.ret_reg = Register(builder.ret_types[-1])
154+
if builder.fn_info.is_generator:
155+
self.ret_reg = builder.make_spill_target(builder.ret_types[-1])
156+
else:
157+
self.ret_reg = Register(builder.ret_types[-1])
158+
# assert needed because of apparent mypy bug... it loses track of the union
159+
# and infers the type as object
160+
assert isinstance(self.ret_reg, (Register, AssignmentTarget))
161+
builder.assign(self.ret_reg, value, line)
156162

157-
builder.add(Assign(self.ret_reg, value))
158163
builder.add(Goto(self.target))
159164

160165

@@ -180,9 +185,8 @@ class FinallyNonlocalControl(CleanupNonlocalControl):
180185
leave and the return register is decrefed if it isn't null.
181186
"""
182187

183-
def __init__(self, outer: NonlocalControl, ret_reg: Optional[Value], saved: Value) -> None:
188+
def __init__(self, outer: NonlocalControl, saved: Value) -> None:
184189
super().__init__(outer)
185-
self.ret_reg = ret_reg
186190
self.saved = saved
187191

188192
def gen_cleanup(self, builder: IRBuilder, line: int) -> None:

0 commit comments

Comments
 (0)