handle multiple cancellations

Author: jakkdl

src/trio/_channel.py

@@ -19,10 +19,11 @@
 from ._abc import ReceiveChannel, ReceiveType, SendChannel, SendType, T
 from ._core import Abort, RaiseCancelT, Task, enable_ki_protection
 from ._util import (
+    MultipleExceptionError,
     NoPublicConstructor,
     final,
     generic_function,
-    raise_saving_context,
+    raise_single_exception_from_group,
 )
 
 if sys.version_info < (3, 11):
@@ -548,15 +549,15 @@ async def context_manager(
                 # abandoned generator if it's still alive.
                 nursery.cancel_scope.cancel()
         except BaseExceptionGroup as eg:
-            first, *rest = eg.exceptions
-            if rest:
+            try:
+                raise_single_exception_from_group(eg)
+            except MultipleExceptionError:
                 # In case user has except* we make it possible for them to handle the
                 # exceptions.
                 raise BaseExceptionGroup(
                     "Encountered exception during cleanup of generator object, as well as exception in the contextmanager body - unable to unwrap.",
                     [eg],
                 ) from None
-            raise_saving_context(first)
 
     async def _move_elems_to_channel(
         agen: AsyncGenerator[T, None],

src/trio/_tests/test_channel.py

@@ -506,19 +506,18 @@ async def agen() -> AsyncGenerator[int]:
 
 
 async def test_as_safe_channel_genexit_finally() -> None:
-    events: list[str] = []
-
     @as_safe_channel
-    async def agen(stuff: list[str]) -> AsyncGenerator[int]:
+    async def agen(events: list[str]) -> AsyncGenerator[int]:
         try:
             yield 1
         except BaseException as e:
-            stuff.append(repr(e))
+            events.append(repr(e))
             raise
         finally:
-            stuff.append("finally")
+            events.append("finally")
             raise ValueError("agen")
 
+    events: list[str] = []
     with RaisesGroup(
         RaisesGroup(
             Matcher(ValueError, match="^agen$"),
@@ -569,7 +568,7 @@ async def test_as_safe_channel_dont_unwrap_user_exceptiongroup() -> None:
     @as_safe_channel
     async def agen() -> AsyncGenerator[None]:
         raise NotImplementedError("not entered")
-        yield
+        yield  # pragma: no cover
 
     with RaisesGroup(Matcher(ValueError, match="bar"), match="foo"):
         async with agen() as _:
@@ -602,22 +601,27 @@ async def handle_value(
 
 async def test_as_safe_channel_multi_cancel() -> None:
     @as_safe_channel
-    async def agen() -> AsyncGenerator[None]:
+    async def agen(events: list[str]) -> AsyncGenerator[None]:
         try:
             yield
         finally:
             # this will give a warning of ASYNC120, although it's not technically a
             # problem of swallowing existing exceptions
-            await trio.lowlevel.checkpoint()
+            try:
+                await trio.lowlevel.checkpoint()
+            except trio.Cancelled:
+                events.append("agen cancel")
+                raise
 
+    events: list[str] = []
     with trio.CancelScope() as cs:
-        with RaisesGroup(
-            RaisesGroup(
-                trio.Cancelled, trio.Cancelled, match="^Exceptions from Trio nursery$"
-            ),
-            match=r"^Encountered exception during cleanup of generator object, as well as exception in the contextmanager body - unable to unwrap.$",
-        ):
-            async with agen() as recv_chan:
+        with pytest.raises(trio.Cancelled):
+            async with agen(events) as recv_chan:
                 async for _ in recv_chan:  # pragma: no branch
                     cs.cancel()
-                    await trio.lowlevel.checkpoint()
+                    try:
+                        await trio.lowlevel.checkpoint()
+                    except trio.Cancelled:
+                        events.append("body cancel")
+                        raise
+    assert events == ["body cancel", "agen cancel"]

src/trio/_tests/test_util.py

@@ -18,6 +18,7 @@
 )
 from .._util import (
     ConflictDetector,
+    MultipleExceptionError,
     NoPublicConstructor,
     coroutine_or_error,
     final,
@@ -288,26 +289,28 @@ async def test_raise_single_exception_from_group() -> None:
     assert excinfo.value.__cause__ == cause
     assert excinfo.value.__context__ == context
 
-    with pytest.raises(ValueError, match="foo") as excinfo:
-        raise_single_exception_from_group(
-            ExceptionGroup("", [ExceptionGroup("", [exc])])
-        )
-    assert excinfo.value.__cause__ == cause
-    assert excinfo.value.__context__ == context
+    # only unwraps one layer of exceptiongroup
+    inner_eg = ExceptionGroup("inner eg", [exc])
+    inner_cause = SyntaxError("inner eg cause")
+    inner_context = TypeError("inner eg context")
+    inner_eg.__cause__ = inner_cause
+    inner_eg.__context__ = inner_context
+    with RaisesGroup(Matcher(ValueError, match="^foo$"), match="^inner eg$") as eginfo:
+        raise_single_exception_from_group(ExceptionGroup("", [inner_eg]))
+    assert eginfo.value.__cause__ == inner_cause
+    assert eginfo.value.__context__ == inner_context
 
     with pytest.raises(ValueError, match="foo") as excinfo:
         raise_single_exception_from_group(
-            BaseExceptionGroup(
-                "", [cancelled, BaseExceptionGroup("", [cancelled, exc])]
-            )
+            BaseExceptionGroup("", [cancelled, cancelled, exc])
         )
     assert excinfo.value.__cause__ == cause
     assert excinfo.value.__context__ == context
 
     # multiple non-cancelled
     eg = ExceptionGroup("", [ValueError("foo"), ValueError("bar")])
     with pytest.raises(
-        AssertionError,
+        MultipleExceptionError,
         match=r"^Attempted to unwrap exceptiongroup with multiple non-cancelled exceptions. This is often caused by a bug in the caller.$",
     ) as excinfo:
         raise_single_exception_from_group(eg)
@@ -328,6 +331,20 @@ async def test_raise_single_exception_from_group() -> None:
     assert excinfo.value.__cause__ is eg_ki
     assert excinfo.value.__context__ is None
 
+    # and same for SystemExit
+    systemexit_ki = BaseExceptionGroup(
+        "",
+        [
+            ValueError("foo"),
+            ValueError("bar"),
+            SystemExit("this exc doesn't get reraised"),
+        ],
+    )
+    with pytest.raises(SystemExit, match=r"^$") as excinfo:
+        raise_single_exception_from_group(systemexit_ki)
+    assert excinfo.value.__cause__ is systemexit_ki
+    assert excinfo.value.__context__ is None
+
     # if we only got cancelled, first one is reraised
     with pytest.raises(trio.Cancelled, match=r"^Cancelled$") as excinfo:
         raise_single_exception_from_group(

src/trio/_util.py

@@ -4,7 +4,6 @@
 import collections.abc
 import inspect
 import signal
-import sys
 from abc import ABCMeta
 from collections.abc import Awaitable, Callable, Sequence
 from functools import update_wrapper
@@ -21,19 +20,20 @@
 
 import trio
 
-if sys.version_info < (3, 11):
-    from exceptiongroup import BaseExceptionGroup
-
 # Explicit "Any" is not allowed
 CallT = TypeVar("CallT", bound=Callable[..., Any])  # type: ignore[explicit-any]
 T = TypeVar("T")
 RetT = TypeVar("RetT")
 
 if TYPE_CHECKING:
+    import sys
     from types import AsyncGeneratorType, TracebackType
 
     from typing_extensions import ParamSpec, Self, TypeVarTuple, Unpack
 
+    if sys.version_info < (3, 11):
+        from exceptiongroup import BaseExceptionGroup
+
     ArgsT = ParamSpec("ArgsT")
     PosArgsT = TypeVarTuple("PosArgsT")
 
@@ -372,12 +372,19 @@ def raise_saving_context(exc: BaseException) -> NoReturn:
         del exc, context
 
 
+class MultipleExceptionError(Exception):
+    """Raised by raise_single_exception_from_group if encountering multiple
+    non-cancelled exceptions."""
+
+
 def raise_single_exception_from_group(
     eg: BaseExceptionGroup[BaseException],
 ) -> NoReturn:
     """This function takes an exception group that is assumed to have at most
     one non-cancelled exception, which it reraises as a standalone exception.
 
+    This exception may be an exceptiongroup itself, in which case it will not be unwrapped.
+
     If a :exc:`KeyboardInterrupt` is encountered, a new KeyboardInterrupt is immediately
     raised with the entire group as cause.
 
@@ -389,30 +396,27 @@ def raise_single_exception_from_group(
     If multiple non-cancelled exceptions are encountered, it raises
     :exc:`AssertionError`.
     """
-    cancelled_exceptions = []
-    noncancelled_exceptions = []
-
-    # subgroup/split retains excgroup structure, so we need to manually traverse
-    def _parse_excg(e: BaseException) -> None:
+    # immediately bail out if there's any KI or SystemExit
+    for e in eg.exceptions:
         if isinstance(e, (KeyboardInterrupt, SystemExit)):
-            # immediately bail out
-            raise KeyboardInterrupt from eg
+            raise type(e) from eg
+
+    cancelled_exception: trio.Cancelled | None = None
+    noncancelled_exception: BaseException | None = None
 
+    for e in eg.exceptions:
         if isinstance(e, trio.Cancelled):
-            cancelled_exceptions.append(e)
-        elif isinstance(e, BaseExceptionGroup):
-            for sub_e in e.exceptions:
-                _parse_excg(sub_e)
+            if cancelled_exception is None:
+                cancelled_exception = e
+        elif noncancelled_exception is None:
+            noncancelled_exception = e
         else:
-            noncancelled_exceptions.append(e)
-
-    _parse_excg(eg)
-
-    if len(noncancelled_exceptions) > 1:
-        raise AssertionError(
-            "Attempted to unwrap exceptiongroup with multiple non-cancelled exceptions. This is often caused by a bug in the caller."
-        ) from eg
-    if len(noncancelled_exceptions) == 1:
-        raise_saving_context(noncancelled_exceptions[0])
-    assert cancelled_exceptions, "internal error"
-    raise_saving_context(cancelled_exceptions[0])
+            raise MultipleExceptionError(
+                "Attempted to unwrap exceptiongroup with multiple non-cancelled exceptions. This is often caused by a bug in the caller."
+            ) from eg
+
+    if noncancelled_exception is not None:
+        raise_saving_context(noncancelled_exception)
+
+    assert cancelled_exception is not None, "group can't be empty"
+    raise_saving_context(cancelled_exception)