diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index 5c32d8f1a50c..f780db2249df 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -12,6 +12,7 @@ from collections.abc import Sequence from typing import Callable +import mypy.nodes from mypy.nodes import ( ARG_NAMED, ARG_POS, @@ -101,6 +102,7 @@ get_exc_info_op, get_exc_value_op, keep_propagating_op, + no_err_occurred_op, raise_exception_op, reraise_exception_op, restore_exc_info_op, @@ -679,7 +681,7 @@ def try_finally_resolve_control( def transform_try_finally_stmt( - builder: IRBuilder, try_body: GenFunc, finally_body: GenFunc + builder: IRBuilder, try_body: GenFunc, finally_body: GenFunc, line: int = -1 ) -> None: """Generalized try/finally handling that takes functions to gen the bodies. @@ -715,6 +717,118 @@ def transform_try_finally_stmt( builder.activate_block(out_block) +def transform_try_finally_stmt_async( + builder: IRBuilder, try_body: GenFunc, finally_body: GenFunc, line: int = -1 +) -> None: + """Async-aware try/finally handling for when finally contains await. + + This version uses a modified approach that preserves exceptions across await.""" + + # We need to handle returns properly, so we'll use TryFinallyNonlocalControl + # to track return values, similar to the regular try/finally implementation + + err_handler, main_entry, return_entry, finally_entry = ( + BasicBlock(), + BasicBlock(), + BasicBlock(), + BasicBlock(), + ) + + # Track if we're returning from the try block + control = TryFinallyNonlocalControl(return_entry) + builder.builder.push_error_handler(err_handler) + builder.nonlocal_control.append(control) + builder.goto_and_activate(BasicBlock()) + try_body() + builder.goto(main_entry) + builder.nonlocal_control.pop() + builder.builder.pop_error_handler() + ret_reg = control.ret_reg + + # Normal case - no exception or return + builder.activate_block(main_entry) + builder.goto(finally_entry) + + # Return case + builder.activate_block(return_entry) + builder.goto(finally_entry) + + # Exception case - need to catch to clear the error indicator + builder.activate_block(err_handler) + # Catch the error to clear Python's error indicator + builder.call_c(error_catch_op, [], line) + # We're not going to use old_exc since it won't survive await + # The exception is now in sys.exc_info() + builder.goto(finally_entry) + + # Finally block + builder.activate_block(finally_entry) + + # Execute finally body + finally_body() + + # After finally, we need to handle exceptions carefully: + # 1. If finally raised a new exception, it's in the error indicator - let it propagate + # 2. If finally didn't raise, check if we need to reraise the original from sys.exc_info() + # 3. If there was a return, return that value + # 4. Otherwise, normal exit + + # First, check if there's a current exception in the error indicator + # (this would be from the finally block) + no_current_exc = builder.call_c(no_err_occurred_op, [], line) + finally_raised = BasicBlock() + check_original = BasicBlock() + builder.add(Branch(no_current_exc, check_original, finally_raised, Branch.BOOL)) + + # Finally raised an exception - let it propagate naturally + builder.activate_block(finally_raised) + builder.call_c(keep_propagating_op, [], NO_TRACEBACK_LINE_NO) + builder.add(Unreachable()) + + # No exception from finally, check if we need to handle return or original exception + builder.activate_block(check_original) + + # Check if we have a return value + if ret_reg: + return_block, check_old_exc = BasicBlock(), BasicBlock() + builder.add(Branch(builder.read(ret_reg), check_old_exc, return_block, Branch.IS_ERROR)) + + builder.activate_block(return_block) + builder.nonlocal_control[-1].gen_return(builder, builder.read(ret_reg), -1) + + builder.activate_block(check_old_exc) + + # Check if we need to reraise the original exception from sys.exc_info + exc_info = builder.call_c(get_exc_info_op, [], line) + exc_type = builder.add(TupleGet(exc_info, 0, line)) + + # Check if exc_type is None + none_obj = builder.none_object() + has_exc = builder.binary_op(exc_type, none_obj, "is not", line) + + reraise_block, exit_block = BasicBlock(), BasicBlock() + builder.add(Branch(has_exc, reraise_block, exit_block, Branch.BOOL)) + + # Reraise the original exception + builder.activate_block(reraise_block) + builder.call_c(reraise_exception_op, [], NO_TRACEBACK_LINE_NO) + builder.add(Unreachable()) + + # Normal exit + builder.activate_block(exit_block) + + +# A simple visitor to detect await expressions +class AwaitDetector(mypy.traverser.TraverserVisitor): + def __init__(self) -> None: + super().__init__() + self.has_await = False + + def visit_await_expr(self, o: mypy.nodes.AwaitExpr) -> None: + self.has_await = True + super().visit_await_expr(o) + + def transform_try_stmt(builder: IRBuilder, t: TryStmt) -> None: # Our compilation strategy for try/except/else/finally is to # treat try/except/else and try/finally as separate language @@ -723,6 +837,17 @@ def transform_try_stmt(builder: IRBuilder, t: TryStmt) -> None: # body of a try/finally block. if t.is_star: builder.error("Exception groups and except* cannot be compiled yet", t.line) + + # Check if we're in an async function with a finally block that contains await + use_async_version = False + if t.finally_body and builder.fn_info.is_coroutine: + detector = AwaitDetector() + t.finally_body.accept(detector) + + if detector.has_await: + # Use the async version that handles exceptions correctly + use_async_version = True + if t.finally_body: def transform_try_body() -> None: @@ -733,7 +858,14 @@ def transform_try_body() -> None: body = t.finally_body - transform_try_finally_stmt(builder, transform_try_body, lambda: builder.accept(body)) + if use_async_version: + transform_try_finally_stmt_async( + builder, transform_try_body, lambda: builder.accept(body), t.line + ) + else: + transform_try_finally_stmt( + builder, transform_try_body, lambda: builder.accept(body), t.line + ) else: transform_try_except_stmt(builder, t) @@ -824,6 +956,7 @@ def finally_body() -> None: builder, lambda: transform_try_except(builder, try_body, [(None, None, except_body)], None, line), finally_body, + line, ) diff --git a/mypyc/test-data/run-async.test b/mypyc/test-data/run-async.test index 2dad720f99cd..d1fb68d9f013 100644 --- a/mypyc/test-data/run-async.test +++ b/mypyc/test-data/run-async.test @@ -946,3 +946,214 @@ test_async_with_mixed_return() [file asyncio/__init__.pyi] def run(x: object) -> object: ... + +[case testAsyncTryExceptFinallyAwait] +import asyncio +from testutil import assertRaises + +class TestError(Exception): + pass + +# Test 0: Simplest case - just try/finally with raise and await +async def simple_try_finally_await() -> None: + try: + raise ValueError("simple error") + finally: + await asyncio.sleep(0) + +# Test 1: Raise inside try, catch in except, don't re-raise +async def async_try_except_no_reraise() -> int: + try: + raise ValueError("test error") + return 1 # Never reached + except ValueError: + return 2 # Should return this + finally: + await asyncio.sleep(0) + return 3 # Should not reach this + +# Test 2: Raise inside try, catch in except, re-raise +async def async_try_except_reraise() -> int: + try: + raise ValueError("test error") + return 1 # Never reached + except ValueError: + raise # Re-raise the exception + finally: + await asyncio.sleep(0) + return 2 # Should not reach this + +# Test 3: Raise inside try, catch in except, raise different error +async def async_try_except_raise_different() -> int: + try: + raise ValueError("original error") + return 1 # Never reached + except ValueError: + raise RuntimeError("different error") + finally: + await asyncio.sleep(0) + return 2 # Should not reach this + +# Test 4: Another try/except block inside finally +async def async_try_except_inside_finally() -> int: + try: + raise ValueError("outer error") + return 1 # Never reached + finally: + await asyncio.sleep(0) + try: + raise RuntimeError("inner error") + except RuntimeError: + pass # Catch inner error + return 2 # What happens after finally with inner exception handled? + +# Test 5: Another try/finally block inside finally +async def async_try_finally_inside_finally() -> int: + try: + raise ValueError("outer error") + return 1 # Never reached + finally: + await asyncio.sleep(0) + try: + raise RuntimeError("inner error") + finally: + await asyncio.sleep(0) + return 2 # Should not reach this + +# Control case: No await in finally - should work correctly +async def async_exception_no_await_in_finally() -> None: + """Control case: This works correctly - exception propagates""" + try: + raise TestError("This exception will propagate!") + finally: + pass # No await here + +# Test function with no exception to check normal flow +async def async_no_exception_with_await_in_finally() -> int: + try: + return 1 # Normal return + finally: + await asyncio.sleep(0) + return 2 # Should not reach this + +def test_async_try_except_finally_await() -> None: + # Test 0: Simplest case - just try/finally with exception + # Expected: ValueError propagates + with assertRaises(ValueError): + asyncio.run(simple_try_finally_await()) + + # Test 1: Exception caught, not re-raised + # Expected: return 2 (from except block) + result = asyncio.run(async_try_except_no_reraise()) + assert result == 2, f"Expected 2, got {result}" + + # Test 2: Exception caught and re-raised + # Expected: ValueError propagates + with assertRaises(ValueError): + asyncio.run(async_try_except_reraise()) + + # Test 3: Exception caught, different exception raised + # Expected: RuntimeError propagates + with assertRaises(RuntimeError): + asyncio.run(async_try_except_raise_different()) + + # Test 4: Try/except inside finally + # Expected: ValueError propagates (outer exception) + with assertRaises(ValueError): + asyncio.run(async_try_except_inside_finally()) + + # Test 5: Try/finally inside finally + # Expected: RuntimeError propagates (inner error) + with assertRaises(RuntimeError): + asyncio.run(async_try_finally_inside_finally()) + + # Control case: No await in finally (should work correctly) + with assertRaises(TestError): + asyncio.run(async_exception_no_await_in_finally()) + + # Test normal flow (no exception) + # Expected: return 1 + result = asyncio.run(async_no_exception_with_await_in_finally()) + assert result == 1, f"Expected 1, got {result}" + +[file asyncio/__init__.pyi] +async def sleep(t: float) -> None: ... +def run(x: object) -> object: ... + +[case testAsyncContextManagerExceptionHandling] +import asyncio +from typing import Optional, Type +from testutil import assertRaises + +# Test 1: Basic async context manager that doesn't suppress exceptions +class AsyncContextManager: + async def __aenter__(self) -> 'AsyncContextManager': + return self + + async def __aexit__(self, exc_type: Optional[Type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: object) -> None: + # This await in __aexit__ is like await in finally + await asyncio.sleep(0) + # Don't suppress the exception (return None/False) + +async def func_with_async_context_manager() -> str: + async with AsyncContextManager(): + raise ValueError("Exception inside async with") + return "should not reach" # Never reached + return "should not reach either" # Never reached + +async def test_basic_exception() -> str: + try: + await func_with_async_context_manager() + return "func_a returned normally - bug!" + except ValueError: + return "caught ValueError - correct!" + except Exception as e: + return f"caught different exception: {type(e).__name__}" + +# Test 2: Async context manager that raises a different exception in __aexit__ +class AsyncContextManagerRaisesInExit: + async def __aenter__(self) -> 'AsyncContextManagerRaisesInExit': + return self + + async def __aexit__(self, exc_type: Optional[Type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: object) -> None: + # This await in __aexit__ is like await in finally + await asyncio.sleep(0) + # Raise a different exception - this should replace the original exception + raise RuntimeError("Exception in __aexit__") + +async def func_with_raising_context_manager() -> str: + async with AsyncContextManagerRaisesInExit(): + raise ValueError("Original exception") + return "should not reach" # Never reached + return "should not reach either" # Never reached + +async def test_exception_in_aexit() -> str: + try: + await func_with_raising_context_manager() + return "func returned normally - unexpected!" + except RuntimeError: + return "caught RuntimeError - correct!" + except ValueError: + return "caught ValueError - original exception not replaced!" + except Exception as e: + return f"caught different exception: {type(e).__name__}" + +def test_async_context_manager_exception_handling() -> None: + # Test 1: Basic exception propagation + result = asyncio.run(test_basic_exception()) + # Expected: "caught ValueError - correct!" + assert result == "caught ValueError - correct!", f"Expected exception to propagate, got: {result}" + + # Test 2: Exception raised in __aexit__ replaces original exception + result = asyncio.run(test_exception_in_aexit()) + # Expected: "caught RuntimeError - correct!" + # (The RuntimeError from __aexit__ should replace the ValueError) + assert result == "caught RuntimeError - correct!", f"Expected RuntimeError from __aexit__, got: {result}" + +[file asyncio/__init__.pyi] +async def sleep(t: float) -> None: ... +def run(x: object) -> object: ...