Skip to content

Commit 09ba1f6

Browse files
Chainfireesarp
authored andcommitted
[mypyc] Fix exception swallowing in async try/finally blocks with await (#19353)
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.
1 parent 5c65e33 commit 09ba1f6

File tree

2 files changed

+346
-2
lines changed

2 files changed

+346
-2
lines changed

mypyc/irbuild/statement.py

Lines changed: 135 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from collections.abc import Sequence
1313
from typing import Callable
1414

15+
import mypy.nodes
1516
from mypy.nodes import (
1617
ARG_NAMED,
1718
ARG_POS,
@@ -101,6 +102,7 @@
101102
get_exc_info_op,
102103
get_exc_value_op,
103104
keep_propagating_op,
105+
no_err_occurred_op,
104106
raise_exception_op,
105107
reraise_exception_op,
106108
restore_exc_info_op,
@@ -679,7 +681,7 @@ def try_finally_resolve_control(
679681

680682

681683
def transform_try_finally_stmt(
682-
builder: IRBuilder, try_body: GenFunc, finally_body: GenFunc
684+
builder: IRBuilder, try_body: GenFunc, finally_body: GenFunc, line: int = -1
683685
) -> None:
684686
"""Generalized try/finally handling that takes functions to gen the bodies.
685687
@@ -715,6 +717,118 @@ def transform_try_finally_stmt(
715717
builder.activate_block(out_block)
716718

717719

720+
def transform_try_finally_stmt_async(
721+
builder: IRBuilder, try_body: GenFunc, finally_body: GenFunc, line: int = -1
722+
) -> None:
723+
"""Async-aware try/finally handling for when finally contains await.
724+
725+
This version uses a modified approach that preserves exceptions across await."""
726+
727+
# We need to handle returns properly, so we'll use TryFinallyNonlocalControl
728+
# to track return values, similar to the regular try/finally implementation
729+
730+
err_handler, main_entry, return_entry, finally_entry = (
731+
BasicBlock(),
732+
BasicBlock(),
733+
BasicBlock(),
734+
BasicBlock(),
735+
)
736+
737+
# Track if we're returning from the try block
738+
control = TryFinallyNonlocalControl(return_entry)
739+
builder.builder.push_error_handler(err_handler)
740+
builder.nonlocal_control.append(control)
741+
builder.goto_and_activate(BasicBlock())
742+
try_body()
743+
builder.goto(main_entry)
744+
builder.nonlocal_control.pop()
745+
builder.builder.pop_error_handler()
746+
ret_reg = control.ret_reg
747+
748+
# Normal case - no exception or return
749+
builder.activate_block(main_entry)
750+
builder.goto(finally_entry)
751+
752+
# Return case
753+
builder.activate_block(return_entry)
754+
builder.goto(finally_entry)
755+
756+
# Exception case - need to catch to clear the error indicator
757+
builder.activate_block(err_handler)
758+
# Catch the error to clear Python's error indicator
759+
builder.call_c(error_catch_op, [], line)
760+
# We're not going to use old_exc since it won't survive await
761+
# The exception is now in sys.exc_info()
762+
builder.goto(finally_entry)
763+
764+
# Finally block
765+
builder.activate_block(finally_entry)
766+
767+
# Execute finally body
768+
finally_body()
769+
770+
# After finally, we need to handle exceptions carefully:
771+
# 1. If finally raised a new exception, it's in the error indicator - let it propagate
772+
# 2. If finally didn't raise, check if we need to reraise the original from sys.exc_info()
773+
# 3. If there was a return, return that value
774+
# 4. Otherwise, normal exit
775+
776+
# First, check if there's a current exception in the error indicator
777+
# (this would be from the finally block)
778+
no_current_exc = builder.call_c(no_err_occurred_op, [], line)
779+
finally_raised = BasicBlock()
780+
check_original = BasicBlock()
781+
builder.add(Branch(no_current_exc, check_original, finally_raised, Branch.BOOL))
782+
783+
# Finally raised an exception - let it propagate naturally
784+
builder.activate_block(finally_raised)
785+
builder.call_c(keep_propagating_op, [], NO_TRACEBACK_LINE_NO)
786+
builder.add(Unreachable())
787+
788+
# No exception from finally, check if we need to handle return or original exception
789+
builder.activate_block(check_original)
790+
791+
# Check if we have a return value
792+
if ret_reg:
793+
return_block, check_old_exc = BasicBlock(), BasicBlock()
794+
builder.add(Branch(builder.read(ret_reg), check_old_exc, return_block, Branch.IS_ERROR))
795+
796+
builder.activate_block(return_block)
797+
builder.nonlocal_control[-1].gen_return(builder, builder.read(ret_reg), -1)
798+
799+
builder.activate_block(check_old_exc)
800+
801+
# Check if we need to reraise the original exception from sys.exc_info
802+
exc_info = builder.call_c(get_exc_info_op, [], line)
803+
exc_type = builder.add(TupleGet(exc_info, 0, line))
804+
805+
# Check if exc_type is None
806+
none_obj = builder.none_object()
807+
has_exc = builder.binary_op(exc_type, none_obj, "is not", line)
808+
809+
reraise_block, exit_block = BasicBlock(), BasicBlock()
810+
builder.add(Branch(has_exc, reraise_block, exit_block, Branch.BOOL))
811+
812+
# Reraise the original exception
813+
builder.activate_block(reraise_block)
814+
builder.call_c(reraise_exception_op, [], NO_TRACEBACK_LINE_NO)
815+
builder.add(Unreachable())
816+
817+
# Normal exit
818+
builder.activate_block(exit_block)
819+
820+
821+
# A simple visitor to detect await expressions
822+
class AwaitDetector(mypy.traverser.TraverserVisitor):
823+
def __init__(self) -> None:
824+
super().__init__()
825+
self.has_await = False
826+
827+
def visit_await_expr(self, o: mypy.nodes.AwaitExpr) -> None:
828+
self.has_await = True
829+
super().visit_await_expr(o)
830+
831+
718832
def transform_try_stmt(builder: IRBuilder, t: TryStmt) -> None:
719833
# Our compilation strategy for try/except/else/finally is to
720834
# treat try/except/else and try/finally as separate language
@@ -723,6 +837,17 @@ def transform_try_stmt(builder: IRBuilder, t: TryStmt) -> None:
723837
# body of a try/finally block.
724838
if t.is_star:
725839
builder.error("Exception groups and except* cannot be compiled yet", t.line)
840+
841+
# Check if we're in an async function with a finally block that contains await
842+
use_async_version = False
843+
if t.finally_body and builder.fn_info.is_coroutine:
844+
detector = AwaitDetector()
845+
t.finally_body.accept(detector)
846+
847+
if detector.has_await:
848+
# Use the async version that handles exceptions correctly
849+
use_async_version = True
850+
726851
if t.finally_body:
727852

728853
def transform_try_body() -> None:
@@ -733,7 +858,14 @@ def transform_try_body() -> None:
733858

734859
body = t.finally_body
735860

736-
transform_try_finally_stmt(builder, transform_try_body, lambda: builder.accept(body))
861+
if use_async_version:
862+
transform_try_finally_stmt_async(
863+
builder, transform_try_body, lambda: builder.accept(body), t.line
864+
)
865+
else:
866+
transform_try_finally_stmt(
867+
builder, transform_try_body, lambda: builder.accept(body), t.line
868+
)
737869
else:
738870
transform_try_except_stmt(builder, t)
739871

@@ -824,6 +956,7 @@ def finally_body() -> None:
824956
builder,
825957
lambda: transform_try_except(builder, try_body, [(None, None, except_body)], None, line),
826958
finally_body,
959+
line,
827960
)
828961

829962

0 commit comments

Comments
 (0)