do everything but unwrapping the exception from inside the group

Author: jakkdl

src/trio/_channel.py

@@ -458,16 +458,16 @@ async def aclose(self) -> None:
 
 class RecvChanWrapper(ReceiveChannel[T]):
     def __init__(
-        self, recv_chan: MemoryReceiveChannel[T], send_semaphore: trio.Semaphore | None
+        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?
+    # TODO: should this allow clones? We'd signal that by inheriting from
+    # MemoryReceiveChannel.
 
     async def receive(self) -> T:
-        if self.send_semaphore is not None:
-            self.send_semaphore.release()
+        self.send_semaphore.release()
         return await self.recv_chan.receive()
 
     async def aclose(self) -> None:
@@ -485,12 +485,9 @@ def __exit__(
         self.recv_chan.close()
 
 
-def background_with_channel(max_buffer_size: int | None = 0) -> Callable[
-    [
-        Callable[P, AsyncGenerator[T, None]],
-    ],
-    Callable[P, AbstractAsyncContextManager[trio.abc.ReceiveChannel[T]]],
-]:
+def background_with_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
@@ -511,7 +508,7 @@ def background_with_channel(max_buffer_size: int | None = 0) -> Callable[
     offering only the safe interface, and you can still write your iterables
     with the convenience of ``yield``.  For example::
 
-        @background_with_channel()
+        @background_with_channel
         async def my_async_iterable(arg, *, kwarg=True):
             while ...:
                 item = await ...
@@ -531,46 +528,30 @@ async def my_async_iterable(arg, *, kwarg=True):
     # Perhaps a future PEP will adopt `async with for` syntax, like
     # https://coconut.readthedocs.io/en/master/DOCS.html#async-with-for
 
-    if not isinstance(max_buffer_size, int) and max_buffer_size is not None:
-        raise TypeError(
-            "`max_buffer_size` must be int or None, not {type(max_buffer_size)}. "
-            "Did you forget the parentheses in `@background_with_channel()`?"
-        )
-
-    def decorator(
-        fn: Callable[P, AsyncGenerator[T, None]],
-    ) -> Callable[P, AbstractAsyncContextManager[trio._channel.RecvChanWrapper[T]]]:
-        @asynccontextmanager
-        @wraps(fn)
-        async def context_manager(
-            *args: P.args, **kwargs: P.kwargs
-        ) -> AsyncGenerator[trio._channel.RecvChanWrapper[T], None]:
-            max_buf_size_float = inf if max_buffer_size is None else max_buffer_size
-            send_chan, recv_chan = trio.open_memory_channel[T](max_buf_size_float)
-            async with trio.open_nursery(strict_exception_groups=True) as nursery:
-                agen = fn(*args, **kwargs)
-                send_semaphore = (
-                    None if max_buffer_size is None else trio.Semaphore(max_buffer_size)
-                )
-                # `nursery.start` to make sure that we will clean up send_chan & agen
-                # If this errors we don't close `recv_chan`, but the caller
-                # never gets access to it, so that's not a problem.
-                await nursery.start(
-                    _move_elems_to_channel, agen, send_chan, send_semaphore
-                )
-                # `async with recv_chan` could eat exceptions, so use sync cm
-                with RecvChanWrapper(recv_chan, send_semaphore) as wrapped_recv_chan:
-                    yield wrapped_recv_chan
-                # User has exited context manager, cancel to immediately close the
-                # abandoned generator if it's still alive.
-                nursery.cancel_scope.cancel()
-
-        return context_manager
+    @asynccontextmanager
+    @wraps(fn)
+    async def context_manager(
+        *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncGenerator[trio._channel.RecvChanWrapper[T], None]:
+        send_chan, recv_chan = trio.open_memory_channel[T](0)
+        async with trio.open_nursery(strict_exception_groups=True) as nursery:
+            agen = fn(*args, **kwargs)
+            send_semaphore = trio.Semaphore(0)
+            # `nursery.start` to make sure that we will clean up send_chan & agen
+            # If this errors we don't close `recv_chan`, but the caller
+            # never gets access to it, so that's not a problem.
+            await nursery.start(_move_elems_to_channel, agen, send_chan, send_semaphore)
+            # `async with recv_chan` could eat exceptions, so use sync cm
+            with RecvChanWrapper(recv_chan, send_semaphore) as wrapped_recv_chan:
+                yield wrapped_recv_chan
+            # User has exited context manager, cancel to immediately close the
+            # abandoned generator if it's still alive.
+            nursery.cancel_scope.cancel()
 
     async def _move_elems_to_channel(
         agen: AsyncGenerator[T, None],
         send_chan: trio.MemorySendChannel[T],
-        send_semaphore: trio.Semaphore | None,
+        send_semaphore: trio.Semaphore,
         task_status: trio.TaskStatus,
     ) -> None:
         # `async with send_chan` will eat exceptions,
@@ -579,22 +560,16 @@ async def _move_elems_to_channel(
             try:
                 task_status.started()
                 while True:
-                    # wait for send_chan to be unblocked
-                    if send_semaphore is not None:
-                        await send_semaphore.acquire()
+                    # wait for receiver to call next on the aiter
+                    await send_semaphore.acquire()
                     try:
                         value = await agen.__anext__()
                     except StopAsyncIteration:
                         return
-                    try:
-                        # Send the value to the channel
-                        await send_chan.send(value)
-                    except trio.BrokenResourceError:
-                        # Closing the corresponding receive channel should cause
-                        # a clean shutdown of the generator.
-                        return
+                    # Send the value to the channel
+                    await send_chan.send(value)
             finally:
                 # replace try-finally with contextlib.aclosing once python39 is dropped
                 await agen.aclose()
 
-    return decorator
+    return context_manager

src/trio/_tests/test_channel.py

@@ -417,7 +417,7 @@ async def do_send(s: trio.MemorySendChannel[int], v: int) -> None:
 
 
 async def test_background_with_channel() -> None:
-    @background_with_channel(1)
+    @background_with_channel
     async def agen() -> AsyncGenerator[int]:
         yield 1
         await trio.sleep_forever()  # simulate deadlock
@@ -429,7 +429,7 @@ async def agen() -> AsyncGenerator[int]:
 
 
 async def test_background_with_channel_exhaust() -> None:
-    @background_with_channel()
+    @background_with_channel
     async def agen() -> AsyncGenerator[int]:
         yield 1
 
@@ -439,7 +439,7 @@ async def agen() -> AsyncGenerator[int]:
 
 
 async def test_background_with_channel_broken_resource() -> None:
-    @background_with_channel()
+    @background_with_channel
     async def agen() -> AsyncGenerator[int]:
         yield 1
         yield 2
@@ -460,7 +460,7 @@ async def agen() -> AsyncGenerator[int]:
 async def test_background_with_channel_cancelled() -> None:
     with trio.CancelScope() as cs:
 
-        @background_with_channel()
+        @background_with_channel
         async def agen() -> AsyncGenerator[None]:  # pragma: no cover
             raise AssertionError(
                 "cancel before consumption means generator should not be iterated"
@@ -476,7 +476,7 @@ async def test_background_with_channel_recv_closed(
 ) -> None:
     event = trio.Event()
 
-    @background_with_channel(1)
+    @background_with_channel
     async def agen() -> AsyncGenerator[int]:
         await event.wait()
         yield 1
@@ -491,7 +491,7 @@ async def agen() -> AsyncGenerator[int]:
 async def test_background_with_channel_no_race() -> None:
     # this previously led to a race condition due to
     # https://github.com/python-trio/trio/issues/1559
-    @background_with_channel()
+    @background_with_channel
     async def agen() -> AsyncGenerator[int]:
         yield 1
         raise ValueError("oae")
@@ -505,7 +505,7 @@ async def agen() -> AsyncGenerator[int]:
 async def test_background_with_channel_buffer_size_too_small(
     autojump_clock: trio.testing.MockClock,
 ) -> None:
-    @background_with_channel(0)
+    @background_with_channel
     async def agen() -> AsyncGenerator[int]:
         yield 1
         raise AssertionError(
@@ -519,27 +519,8 @@ async def agen() -> AsyncGenerator[int]:
                 await trio.sleep_forever()
 
 
-async def test_background_with_channel_buffer_size_just_right(
-    autojump_clock: trio.testing.MockClock,
-) -> None:
-    event = trio.Event()
-
-    @background_with_channel(2)
-    async def agen() -> AsyncGenerator[int]:
-        yield 1
-        event.set()
-        yield 2
-
-    async with agen() as recv_chan:
-        await event.wait()
-        assert await recv_chan.__anext__() == 1
-        assert await recv_chan.__anext__() == 2
-        with pytest.raises(StopAsyncIteration):
-            await recv_chan.__anext__()
-
-
 async def test_background_with_channel_no_interleave() -> None:
-    @background_with_channel()
+    @background_with_channel
     async def agen() -> AsyncGenerator[int]:
         yield 1
         raise AssertionError  # pragma: no cover
@@ -549,30 +530,10 @@ async def agen() -> AsyncGenerator[int]:
         await trio.lowlevel.checkpoint()
 
 
-async def test_background_with_channel_multiple_errors() -> None:
-    event = trio.Event()
-
-    @background_with_channel(1)
-    async def agen() -> AsyncGenerator[int]:
-        yield 1
-        event.set()
-        raise ValueError("agen")
-
-    with RaisesGroup(
-        Matcher(ValueError, match="^agen$"),
-        Matcher(TypeError, match="^iterator$"),
-    ):
-        async with agen() as recv_chan:
-            async for i in recv_chan:  # pragma: no branch
-                assert i == 1
-                await event.wait()
-                raise TypeError("iterator")
-
-
 async def test_background_with_channel_genexit_finally() -> None:
     events: list[str] = []
 
-    @background_with_channel()
+    @background_with_channel
     async def agen(stuff: list[str]) -> AsyncGenerator[int]:
         try:
             yield 1
@@ -596,7 +557,7 @@ async def agen(stuff: list[str]) -> AsyncGenerator[int]:
 
 
 async def test_background_with_channel_nested_loop() -> None:
-    @background_with_channel()
+    @background_with_channel
     async def agen() -> AsyncGenerator[int]:
         for i in range(2):
             yield i
@@ -610,34 +571,3 @@ async def agen() -> AsyncGenerator[int]:
                     assert (i, j) == (ii, jj)
                     jj += 1
             ii += 1
-
-
-async def test_background_with_channel_no_parens() -> None:
-    with pytest.raises(TypeError, match="must be int or None"):
-
-        @background_with_channel  # type: ignore[arg-type]
-        async def agen() -> AsyncGenerator[None]:
-            yield  # pragma: no cover
-
-
-async def test_background_with_channel_inf_buffer() -> None:
-    event = trio.Event()
-
-    # agen immediately starts yielding numbers
-    # into the buffer upon entering the cm
-    @background_with_channel(None)
-    async def agen() -> AsyncGenerator[int]:
-        for i in range(10):
-            yield i
-        event.set()
-        # keep agen alive to receive values
-        await trio.sleep_forever()
-
-    async with agen() as recv_chan:
-        await event.wait()
-        j = 0
-        async for i in recv_chan:
-            assert i == j
-            j += 1
-            if j == 10:
-                break