diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index 21f47b190c30..d978265339c4 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -1351,30 +1351,46 @@ def emit_yield_from_or_await( stop_block, main_block, done_block = BasicBlock(), BasicBlock(), BasicBlock() - if isinstance(iter_reg.type, RInstance) and iter_reg.type.class_ir.has_method(helper_method): - # Second fast path optimization: call helper directly (see also comment above). - # - # Calling a generated generator, so avoid raising StopIteration by passing - # an extra PyObject ** argument to helper where the stop iteration value is stored. - fast_path = True + fast_path = isinstance(iter_reg.type, RInstance) and iter_reg.type.class_ir.has_method( + helper_method + ) + + # Register where a native child stores its return value instead of raising + # StopIteration (only used on the fast path). + stop_iter_val = Register(object_rprimitive) if fast_path else None + + def native_step(sent: Value) -> Value: + """Advance a native generator/coroutine child by calling its helper directly. + + Returns the value yielded by the child, or NULL if the child completed or + raised. On normal completion the return value is stored in stop_iter_val, + which is set to the error value if a real exception was raised instead. + """ + assert stop_iter_val is not None obj = builder.read(iter_reg, line) nn = builder.none_object() - stop_iter_val = Register(object_rprimitive) err = builder.add(LoadErrorValue(object_rprimitive, undefines=True)) builder.assign(stop_iter_val, err, line) ptr = builder.add(LoadAddress(object_pointer_rprimitive, stop_iter_val)) - m = MethodCall(obj, helper_method, [nn, nn, nn, nn, ptr], line) + m = MethodCall(obj, helper_method, [nn, nn, nn, sent, ptr], line) # Generators have custom error handling, so disable normal error handling. m.error_kind = ERR_NEVER - _y_init = builder.add(m) + return builder.add(m) + + if fast_path: + # Second fast path optimization: call helper directly (see also comment above). + # + # Calling a generated generator, so avoid raising StopIteration by passing + # an extra PyObject ** argument to helper where the stop iteration value is stored. + _y_init = native_step(builder.none_object()) else: - fast_path = False _y_init = builder.call_c(next_raw_op, [builder.read(iter_reg, line)], line) builder.add(Branch(_y_init, stop_block, main_block, Branch.IS_ERROR)) builder.activate_block(stop_block) if fast_path: + assert stop_iter_val is not None builder.primitive_op(propagate_if_error_op, [stop_iter_val], line) builder.assign(result, stop_iter_val, line) else: @@ -1423,11 +1439,18 @@ def except_body() -> None: builder.nonlocal_control[-1].gen_break(builder, 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, line), builder.read(received_reg, line)], line - ) + # This path runs when the parent's yield is resumed normally via next() or send(). + # An exception injected via throw() or close() takes the except_body path instead. + if fast_path: + # Reuse the direct helper call on resumes as well, so that native-to-native + # completion doesn't have to go through .send() and StopIteration. + _y = native_step(builder.read(received_reg, line)) + else: + # 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, line), builder.read(received_reg, line)], line + ) ok, stop = BasicBlock(), BasicBlock() builder.add(Branch(_y, stop, ok, Branch.IS_ERROR)) @@ -1436,10 +1459,17 @@ def else_body() -> None: builder.assign(to_yield_reg, _y, line) builder.nonlocal_control[-1].gen_continue(builder, 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, [], line), line) + if fast_path: + assert stop_iter_val is not None + # The child either returned a value through the out pointer, or raised + # a real exception (in which case this propagates it). + builder.primitive_op(propagate_if_error_op, [stop_iter_val], line) + builder.assign(result, stop_iter_val, line) + else: + # Try extracting a return value from a StopIteration and return it. + # If it wasn't, this rereaises the exception. + builder.assign(result, builder.call_c(check_stop_op, [], line), line) builder.nonlocal_control[-1].gen_break(builder, line) builder.push_loop_stack(loop_block, done_block) diff --git a/mypyc/test-data/run-async.test b/mypyc/test-data/run-async.test index 7faa3a2478b5..bdffc658486a 100644 --- a/mypyc/test-data/run-async.test +++ b/mypyc/test-data/run-async.test @@ -2409,3 +2409,92 @@ async def test_borrow_final_attr_across_await_after_comprehension() -> None: async def sleep(t: float) -> None: ... [typing fixtures/typing-full.pyi] + +[case testRunAsyncResumedNativeAwait] +from typing import Any, Generator + +from testutil import assertRaises + +class MyError(Exception): + pass + +class Suspend: + """Awaitable that suspends the given number of times before returning.""" + + def __init__(self, n: int) -> None: + self.n = n + + def __await__(self) -> Generator[Any, Any, int]: + i = 0 + while i < self.n: + yield i + i += 1 + return self.n + +async def child(n: int) -> int: + total = await Suspend(n) + return total + 1 + +async def parent(n: int) -> int: + # Statically known native child that suspends before completing, so the + # await loop must resume it after a suspension. + return await child(n) + 10 + +async def grandparent(n: int) -> int: + return await parent(n) + 100 + +async def tuple_child(n: int) -> tuple[int, str]: + await Suspend(n) + return (n, "x") + +async def tuple_parent(n: int) -> tuple[int, str]: + return await tuple_child(n) + +async def raising_child(n: int) -> int: + await Suspend(n) + raise MyError() + +async def raising_parent(n: int) -> int: + return await raising_child(n) + +async def catching_parent(n: int) -> int: + try: + return await raising_child(n) + except MyError: + return -1 + +def drive(coro: Any) -> Any: + """Drive a coroutine to completion without an event loop.""" + steps = 0 + while True: + try: + coro.send(None) + except StopIteration as e: + return e.value + steps += 1 + assert steps < 100 + +def test_resumed_native_await() -> None: + for n in range(5): + assert drive(parent(n)) == n + 11 + assert drive(grandparent(n)) == n + 111 + assert drive(tuple_parent(n)) == (n, "x") + +def test_resumed_native_await_exception() -> None: + for n in range(5): + with assertRaises(MyError): + drive(raising_parent(n)) + assert drive(catching_parent(n)) == -1 + +def test_throw_into_suspended_native_child() -> None: + coro = parent(3) + coro.send(None) + with assertRaises(MyError): + coro.throw(MyError()) + +def test_close_suspended_native_child() -> None: + coro = parent(3) + coro.send(None) + coro.close() + +[typing fixtures/typing-full.pyi]