[mypyc] Fix exception swallowing in async try/finally blocks with await

When a try/finally block in an async function contains an await statement
in the finally block, exceptions raised in the try block are silently
swallowed if a context switch occurs. This happens because mypyc stores
exception information in registers that don't survive across await points.

The Problem:
- mypyc's transform_try_finally_stmt uses error_catch_op to save exceptions
  to a register, then reraise_exception_op to restore from that register
- When await causes a context switch, register values are lost
- The exception information is gone, causing silent exception swallowing

The Solution:
- Add new transform_try_finally_stmt_async for async-aware exception handling
- Use sys.exc_info() to preserve exceptions across context switches instead
  of registers
- Check error indicator first to handle new exceptions raised in finally
- Route to async version when finally block contains await expressions

Implementation Details:
- transform_try_finally_stmt_async uses get_exc_info_op/restore_exc_info_op
  which work with sys.exc_info() that survives context switches
- Proper exception priority: new exceptions in finally replace originals
- Added has_await_in_block helper to detect await expressions

Test Coverage:
Added comprehensive async exception handling tests:
- testAsyncTryExceptFinallyAwait: 8 test cases covering various scenarios
  - Simple try/finally with exception and await
  - Exception caught but not re-raised
  - Exception caught and re-raised
  - Different exception raised in except
  - Try/except inside finally block
  - Try/finally inside finally block
  - Control case without await
  - Normal flow without exceptions
- testAsyncContextManagerExceptionHandling: Verifies async with still works
  - Basic exception propagation
  - Exception in __aexit__ replacing original

  See mypyc/mypyc#1114

Author: Chainfire

mypy/errorcodes.py

@@ -282,13 +282,6 @@ def __hash__(self) -> int:
 # This is a catch-all for remaining uncategorized errors.
 MISC: Final[ErrorCode] = ErrorCode("misc", "Miscellaneous other checks", "General")
 
-# Mypyc-specific error codes
-MYPYC_TRY_FINALLY_AWAIT: Final[ErrorCode] = ErrorCode(
-    "mypyc-try-finally-await",
-    "Async try/finally blocks with await in finally are not supported by mypyc",
-    "General",
-)
-
 OVERLOAD_CANNOT_MATCH: Final[ErrorCode] = ErrorCode(
     "overload-cannot-match",
     "Warn if an @overload signature can never be matched",

mypyc/irbuild/statement.py

@@ -13,8 +13,6 @@
 from typing import Callable
 
 import mypy.nodes
-import mypy.traverser
-from mypy.errorcodes import MYPYC_TRY_FINALLY_AWAIT
 from mypy.nodes import (
     ARG_NAMED,
     ARG_POS,
@@ -103,6 +101,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,
@@ -712,6 +711,104 @@ 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
+    old_exc = 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:
@@ -733,32 +830,14 @@ def transform_try_stmt(builder: IRBuilder, t: TryStmt) -> None:
         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:
-            # Check if this error is suppressed with # type: ignore
-            error_ignored = False
-            if builder.module_name in builder.graph:
-                mypyfile = builder.graph[builder.module_name].tree
-                if mypyfile and t.line in mypyfile.ignored_lines:
-                    ignored_codes = mypyfile.ignored_lines[t.line]
-                    # Empty list means ignore all errors on this line
-                    # Otherwise check for specific error code
-                    if not ignored_codes or "mypyc-try-finally-await" in ignored_codes:
-                        error_ignored = True
-
-            if not error_ignored:
-                builder.error(
-                    "try/(except/)finally blocks in async functions with 'await' in "
-                    "the finally block are not supported by mypyc. Exceptions "
-                    "(re-)raised in the try or except blocks will be silently "
-                    "swallowed if a context switch occurs. Ignore with "
-                    "'# type: ignore[mypyc-try-finally-await, unused-ignore]', if you "
-                    "really know what you're doing.",
-                    t.line
-                )
+            # Use the async version that handles exceptions correctly
+            use_async_version = True
 
     if t.finally_body:
 
@@ -770,7 +849,10 @@ def transform_try_body() -> None:
 
         body = t.finally_body
 
-        transform_try_finally_stmt(builder, transform_try_body, lambda: builder.accept(body), t.line)
+        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)
 

mypyc/test-data/run-async.test

@@ -645,9 +645,6 @@ def test_async_def_contains_two_nested_functions() -> None:
 def run(x: object) -> object: ...
 
 [case testAsyncTryExceptFinallyAwait]
-# Comprehensive test for bug where exceptions are swallowed in async functions
-# when the finally block contains an await statement
-
 import asyncio
 from testutil import assertRaises
 
@@ -656,14 +653,14 @@ class TestError(Exception):
 
 # Test 0: Simplest case - just try/finally with raise and await
 async def simple_try_finally_await() -> None:
-    try:  # type: ignore[mypyc-try-finally-await]
+    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:  # type: ignore[mypyc-try-finally-await]
+    try:
         raise ValueError("test error")
         return 1  # Never reached
     except ValueError:
@@ -674,7 +671,7 @@ async def async_try_except_no_reraise() -> int:
 
 # Test 2: Raise inside try, catch in except, re-raise
 async def async_try_except_reraise() -> int:
-    try:  # type: ignore[mypyc-try-finally-await]
+    try:
         raise ValueError("test error")
         return 1  # Never reached
     except ValueError:
@@ -685,7 +682,7 @@ async def async_try_except_reraise() -> int:
 
 # Test 3: Raise inside try, catch in except, raise different error
 async def async_try_except_raise_different() -> int:
-    try:  # type: ignore[mypyc-try-finally-await]
+    try:
         raise ValueError("original error")
         return 1  # Never reached
     except ValueError:
@@ -696,7 +693,7 @@ async def async_try_except_raise_different() -> int:
 
 # Test 4: Another try/except block inside finally
 async def async_try_except_inside_finally() -> int:
-    try:  # type: ignore[mypyc-try-finally-await]
+    try:
         raise ValueError("outer error")
         return 1  # Never reached
     finally:
@@ -709,12 +706,12 @@ async def async_try_except_inside_finally() -> int:
 
 # Test 5: Another try/finally block inside finally
 async def async_try_finally_inside_finally() -> int:
-    try:  # type: ignore[mypyc-try-finally-await]
+    try:
         raise ValueError("outer error")
         return 1  # Never reached
     finally:
         await asyncio.sleep(0)
-        try:  # type: ignore[mypyc-try-finally-await]
+        try:
             raise RuntimeError("inner error")
         finally:
             await asyncio.sleep(0)
@@ -730,7 +727,7 @@ async def async_exception_no_await_in_finally() -> None:
 
 # Test function with no exception to check normal flow
 async def async_no_exception_with_await_in_finally() -> int:
-    try:  # type: ignore[mypyc-try-finally-await]
+    try:
         return 1  # Normal return
     finally:
         await asyncio.sleep(0)
@@ -781,20 +778,17 @@ async def sleep(t: float) -> None: ...
 def run(x: object) -> object: ...
 
 [case testAsyncContextManagerExceptionHandling]
-# Test async context managers with exceptions
-# Async context managers use try/finally internally but seem to work
-# correctly
-
 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: type[BaseException] | None,
-                       exc_val: BaseException | None,
+    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)
@@ -820,8 +814,8 @@ class AsyncContextManagerRaisesInExit:
     async def __aenter__(self) -> 'AsyncContextManagerRaisesInExit':
         return self
 
-    async def __aexit__(self, exc_type: type[BaseException] | None,
-                       exc_val: BaseException | None,
+    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)