rename background_with_channel -> as_safe_channel, rename _raise, don't unwrap user exception groups, add test for multiple receivers, clean up docs

Author: jakkdl

docs/source/reference-core.rst

@@ -1609,10 +1609,10 @@ the numbers 0 through 9 with a 1-second delay before each one:
 
 Trio supports async generators, but there's several caveats and it's very
 hard to handle them properly. Therefore Trio bundles a helper,
-`trio.background_with_channel` that does it for you.
+`trio.as_safe_channel` that does it for you.
 
 
-.. autofunction:: trio.background_with_channel
+.. autofunction:: trio.as_safe_channel
 
 The details behind the problems are described in the following sections.
 

newsfragments/3197.feature.rst

@@ -1 +1 @@
-Add :func:`@trio.background_with_channel <trio.background_with_channel>`, a wrapper that can be used to make async generators safe. This will be the suggested fix for `ASYNC900 <https://flake8-async.readthedocs.io/en/latest/rules.html#async900>`_.
+Add :func:`@trio.as_safe_channel <trio.as_safe_channel>`, a wrapper that can be used to make async generators safe. This will be the suggested fix for `ASYNC900 <https://flake8-async.readthedocs.io/en/latest/rules.html#async900>`_.

src/trio/__init__.py

@@ -27,7 +27,7 @@
     MemoryChannelStatistics as MemoryChannelStatistics,
     MemoryReceiveChannel as MemoryReceiveChannel,
     MemorySendChannel as MemorySendChannel,
-    background_with_channel as background_with_channel,
+    as_safe_channel as as_safe_channel,
     open_memory_channel as open_memory_channel,
 )
 from ._core import (

src/trio/_channel.py

@@ -22,7 +22,7 @@
     NoPublicConstructor,
     final,
     generic_function,
-    raise_single_exception_from_group,
+    raise_saving_context,
 )
 
 if sys.version_info < (3, 11):
@@ -468,18 +468,15 @@ class RecvChanWrapper(ReceiveChannel[T]):
     def __init__(
         self, recv_chan: MemoryReceiveChannel[T], send_semaphore: trio.Semaphore
     ) -> None:
-        self.recv_chan = recv_chan
-        self.send_semaphore = send_semaphore
-
-    # TODO: should this allow clones? We'd signal that by inheriting from
-    # MemoryReceiveChannel.
+        self._recv_chan = recv_chan
+        self._send_semaphore = send_semaphore
 
     async def receive(self) -> T:
-        self.send_semaphore.release()
-        return await self.recv_chan.receive()
+        self._send_semaphore.release()
+        return await self._recv_chan.receive()
 
     async def aclose(self) -> None:
-        await self.recv_chan.aclose()
+        await self._recv_chan.aclose()
 
     def __enter__(self) -> Self:
         return self
@@ -490,19 +487,14 @@ def __exit__(
         exc_value: BaseException | None,
         traceback: TracebackType | None,
     ) -> None:
-        self.recv_chan.close()
+        self._recv_chan.close()
 
 
-def background_with_channel(
+def as_safe_channel(
     fn: Callable[P, AsyncGenerator[T, None]],
 ) -> Callable[P, AbstractAsyncContextManager[ReceiveChannel[T]]]:
     """Decorate an async generator function to make it cancellation-safe.
 
-    This is mostly a drop-in replacement, except for the fact that it will
-    wrap errors in exception groups due to the internal nursery. Although when
-    using it without a buffer it should be exceedingly rare to get multiple
-    exceptions.
-
     The ``yield`` keyword offers a very convenient way to write iterators...
     which makes it really unfortunate that async generators are so difficult
     to call correctly.  Yielding from the inside of a cancel scope or a nursery
@@ -516,7 +508,7 @@ def background_with_channel(
     offering only the safe interface, and you can still write your iterables
     with the convenience of ``yield``.  For example::
 
-        @background_with_channel
+        @as_safe_channel
         async def my_async_iterable(arg, *, kwarg=True):
             while ...:
                 item = await ...
@@ -529,9 +521,6 @@ async def my_async_iterable(arg, *, kwarg=True):
     While the combined async-with-async-for can be inconvenient at first,
     the context manager is indispensable for both correctness and for prompt
     cleanup of resources.
-
-    If you specify ``max_buffer_size>0`` the async generator will run concurrently
-    with your iterator, until the buffer is full.
     """
     # Perhaps a future PEP will adopt `async with for` syntax, like
     # https://coconut.readthedocs.io/en/master/DOCS.html#async-with-for
@@ -559,12 +548,15 @@ async def context_manager(
                 # abandoned generator if it's still alive.
                 nursery.cancel_scope.cancel()
         except BaseExceptionGroup as eg:
-            try:
-                raise_single_exception_from_group(eg)
-            except AssertionError:
-                raise RuntimeError(
-                    "Encountered exception during cleanup of generator object, as well as exception in the contextmanager body."
-                ) from eg
+            first, *rest = eg.exceptions
+            if rest:
+                # 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

@@ -1,14 +1,18 @@
 from __future__ import annotations
 
+import sys
 from typing import TYPE_CHECKING, Union
 
 import pytest
 
 import trio
-from trio import EndOfChannel, background_with_channel, open_memory_channel
+from trio import EndOfChannel, as_safe_channel, open_memory_channel
 
 from ..testing import Matcher, RaisesGroup, assert_checkpoints, wait_all_tasks_blocked
 
+if sys.version_info < (3, 11):
+    from exceptiongroup import ExceptionGroup
+
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator
 
@@ -416,8 +420,8 @@ async def do_send(s: trio.MemorySendChannel[int], v: int) -> None:
         r.receive_nowait()
 
 
-async def test_background_with_channel_exhaust() -> None:
-    @background_with_channel
+async def test_as_safe_channel_exhaust() -> None:
+    @as_safe_channel
     async def agen() -> AsyncGenerator[int]:
         yield 1
 
@@ -426,8 +430,8 @@ async def agen() -> AsyncGenerator[int]:
             assert x == 1
 
 
-async def test_background_with_channel_broken_resource() -> None:
-    @background_with_channel
+async def test_as_safe_channel_broken_resource() -> None:
+    @as_safe_channel
     async def agen() -> AsyncGenerator[int]:
         yield 1
         yield 2
@@ -445,10 +449,10 @@ async def agen() -> AsyncGenerator[int]:
         # but we don't get an error on exit of the cm
 
 
-async def test_background_with_channel_cancelled() -> None:
+async def test_as_safe_channel_cancelled() -> None:
     with trio.CancelScope() as cs:
 
-        @background_with_channel
+        @as_safe_channel
         async def agen() -> AsyncGenerator[None]:  # pragma: no cover
             raise AssertionError(
                 "cancel before consumption means generator should not be iterated"
@@ -459,12 +463,12 @@ async def agen() -> AsyncGenerator[None]:  # pragma: no cover
             cs.cancel()
 
 
-async def test_background_with_channel_recv_closed(
+async def test_as_safe_channel_recv_closed(
     autojump_clock: trio.testing.MockClock,
 ) -> None:
     event = trio.Event()
 
-    @background_with_channel
+    @as_safe_channel
     async def agen() -> AsyncGenerator[int]:
         await event.wait()
         yield 1
@@ -476,10 +480,10 @@ async def agen() -> AsyncGenerator[int]:
         await trio.sleep(1)
 
 
-async def test_background_with_channel_no_race() -> None:
+async def test_as_safe_channel_no_race() -> None:
     # this previously led to a race condition due to
     # https://github.com/python-trio/trio/issues/1559
-    @background_with_channel
+    @as_safe_channel
     async def agen() -> AsyncGenerator[int]:
         yield 1
         raise ValueError("oae")
@@ -490,10 +494,10 @@ async def agen() -> AsyncGenerator[int]:
                 assert x == 1
 
 
-async def test_background_with_channel_buffer_size_too_small(
+async def test_as_safe_channel_buffer_size_too_small(
     autojump_clock: trio.testing.MockClock,
 ) -> None:
-    @background_with_channel
+    @as_safe_channel
     async def agen() -> AsyncGenerator[int]:
         yield 1
         raise AssertionError(
@@ -507,8 +511,8 @@ async def agen() -> AsyncGenerator[int]:
                 await trio.sleep_forever()
 
 
-async def test_background_with_channel_no_interleave() -> None:
-    @background_with_channel
+async def test_as_safe_channel_no_interleave() -> None:
+    @as_safe_channel
     async def agen() -> AsyncGenerator[int]:
         yield 1
         raise AssertionError  # pragma: no cover
@@ -518,10 +522,10 @@ async def agen() -> AsyncGenerator[int]:
         await trio.lowlevel.checkpoint()
 
 
-async def test_background_with_channel_genexit_finally() -> None:
+async def test_as_safe_channel_genexit_finally() -> None:
     events: list[str] = []
 
-    @background_with_channel
+    @as_safe_channel
     async def agen(stuff: list[str]) -> AsyncGenerator[int]:
         try:
             yield 1
@@ -532,24 +536,23 @@ async def agen(stuff: list[str]) -> AsyncGenerator[int]:
             stuff.append("finally")
             raise ValueError("agen")
 
-    with pytest.raises(
-        RuntimeError,
-        match=r"^Encountered exception during cleanup of generator object, as well as exception in the contextmanager body.$",
-    ) as excinfo:
+    with RaisesGroup(
+        RaisesGroup(
+            Matcher(ValueError, match="^agen$"),
+            Matcher(TypeError, match="^iterator$"),
+        ),
+        match=r"^Encountered exception during cleanup of generator object, as well as exception in the contextmanager body - unable to unwrap.$",
+    ):
         async with agen(events) as recv_chan:
             async for i in recv_chan:  # pragma: no branch
                 assert i == 1
                 raise TypeError("iterator")
 
     assert events == ["GeneratorExit()", "finally"]
-    RaisesGroup(
-        Matcher(ValueError, match="^agen$"),
-        Matcher(TypeError, match="^iterator$"),
-    ).matches(excinfo.value.__cause__)
 
 
-async def test_background_with_channel_nested_loop() -> None:
-    @background_with_channel
+async def test_as_safe_channel_nested_loop() -> None:
+    @as_safe_channel
     async def agen() -> AsyncGenerator[int]:
         for i in range(2):
             yield i
@@ -565,15 +568,49 @@ async def agen() -> AsyncGenerator[int]:
             ii += 1
 
 
-async def test_doesnt_leak_cancellation() -> None:
-    @background_with_channel
-    async def agenfn() -> AsyncGenerator[None]:
+async def test_as_safe_channel_doesnt_leak_cancellation() -> None:
+    @as_safe_channel
+    async def agen() -> AsyncGenerator[None]:
         with trio.CancelScope() as cscope:
             cscope.cancel()
             yield
 
     with pytest.raises(AssertionError):
-        async with agenfn() as recv_chan:
+        async with agen() as recv_chan:
             async for _ in recv_chan:
                 pass
         raise AssertionError("should be reachable")
+
+
+async def test_as_safe_channel_dont_unwrap_user_exceptiongroup() -> None:
+    @as_safe_channel
+    async def agen() -> AsyncGenerator[None]:
+        yield
+
+    with RaisesGroup(Matcher(ValueError, match="bar"), match="foo"):
+        async with agen() as _:
+            raise ExceptionGroup("foo", [ValueError("bar")])
+
+
+async def test_as_safe_channel_multiple_receiver() -> None:
+    event = trio.Event()
+
+    @as_safe_channel
+    async def agen() -> AsyncGenerator[int]:
+        await event.wait()
+        for i in range(2):
+            yield i
+
+    async def handle_value(
+        recv_chan: trio.abc.ReceiveChannel[int],
+        value: int,
+        task_status: trio.TaskStatus,
+    ) -> None:
+        task_status.started()
+        assert await recv_chan.receive() == value
+
+    async with agen() as recv_chan:
+        async with trio.open_nursery() as nursery:
+            await nursery.start(handle_value, recv_chan, 0)
+            await nursery.start(handle_value, recv_chan, 1)
+            event.set()

src/trio/_util.py

@@ -359,9 +359,10 @@ def wraps(  # type: ignore[explicit-any]
     from functools import wraps  # noqa: F401  # this is re-exported
 
 
-def _raise(exc: BaseException) -> NoReturn:
+def raise_saving_context(exc: BaseException) -> NoReturn:
     """This helper allows re-raising an exception without __context__ being set."""
     # cause does not need special handling, we simply avoid using `raise .. from ..`
+    # __suppress_context__ also does not need handling, it's only set if modifying cause
     __tracebackhide__ = True
     context = exc.__context__
     try:
@@ -412,6 +413,6 @@ def _parse_excg(e: BaseException) -> None:
             "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(noncancelled_exceptions[0])
+        raise_saving_context(noncancelled_exceptions[0])
     assert cancelled_exceptions, "internal error"
-    _raise(cancelled_exceptions[0])
+    raise_saving_context(cancelled_exceptions[0])