[mypyc] Fix broken cancellation handling in async def

mypyc/irbuild/builder.py

@@ -700,7 +700,12 @@ def get_assignment_target(
         assert False, "Unsupported lvalue: %r" % lvalue
 
     def read(
-        self, target: Value | AssignmentTarget, line: int = -1, can_borrow: bool = False
+        self,
+        target: Value | AssignmentTarget,
+        line: int = -1,
+        *,
+        can_borrow: bool = False,
+        allow_error_value: bool = False,
     ) -> Value:
         if isinstance(target, Value):
             return target
@@ -716,7 +721,15 @@ def read(
         if isinstance(target, AssignmentTargetAttr):
             if isinstance(target.obj.type, RInstance) and target.obj.type.class_ir.is_ext_class:
                 borrow = can_borrow and target.can_borrow
-                return self.add(GetAttr(target.obj, target.attr, line, borrow=borrow))
+                return self.add(
+                    GetAttr(
+                        target.obj,
+                        target.attr,
+                        line,
+                        borrow=borrow,
+                        allow_error_value=allow_error_value,
+                    )
+                )
             else:
                 return self.py_get_attr(target.obj, target.attr, line)
 

mypyc/irbuild/statement.py

@@ -829,7 +829,14 @@ def transform_try_finally_stmt_async(
     # 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.add(
+            Branch(
+                builder.read(ret_reg, allow_error_value=True),
+                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)

mypyc/test-data/fixtures/testutil.py

@@ -44,7 +44,7 @@
 def assertRaises(typ: type, msg: str = '') -> Iterator[None]:
     try:
         yield
-    except Exception as e:
+    except BaseException as e:
         assert type(e) is typ, f"{e!r} is not a {typ.__name__}"
         assert msg in str(e), f'Message "{e}" does not match "{msg}"'
     else:

mypyc/test-data/run-async.test

@@ -1234,3 +1234,60 @@ def test_callable_arg_same_name_as_helper() -> None:
 
 [file asyncio/__init__.pyi]
 def run(x: object) -> object: ...
+
+[case testRunAsyncCancelFinallySpecialCase]
+import asyncio
+
+from testutil import assertRaises
+
+# Greatly simplified from asyncio.Condition
+class Condition:
+    async def acquire(self) -> None: pass
+
+    async def wait(self) -> bool:
+         l = asyncio.get_running_loop()
+         fut = l.create_future()
+         a = []
+         try:
+             try:
+                 a.append(fut)
+                 try:
+                     await fut
+                     return True
+                 finally:
+                     a.pop()
+             finally:
+                err = None
+                while True:
+                    try:
+                        await self.acquire()
+                        break
+                    except asyncio.CancelledError as e:
+                        err = e
+
+                if err is not None:
+                    try:
+                        raise err
+                    finally:
+                        err = None
+         except BaseException:
+             raise
+
+async def do_cancel() -> None:
+    cond = Condition()
+    wait = asyncio.create_task(cond.wait())
+    asyncio.get_running_loop().call_soon(wait.cancel)
+    with assertRaises(asyncio.CancelledError):
+        await wait
+
+def test_cancel_special_case() -> None:
+    asyncio.run(do_cancel())
+
+[file asyncio/__init__.pyi]
+from typing import Any
+
+class CancelledError(Exception): ...
+
+def run(x: object) -> object: ...
+def get_running_loop() -> Any: ...
+def create_task(x: object) -> Any: ...