add @background_with_channel

src/trio/__init__.py

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

src/trio/_channel.py

@@ -1,10 +1,15 @@
 from __future__ import annotations
 
+import sys
 from collections import OrderedDict, deque
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+from functools import wraps
 from math import inf
 from typing import (
     TYPE_CHECKING,
     Generic,
+    Protocol,
+    TypeVar,
 )
 
 import attrs
@@ -17,9 +22,31 @@
 from ._util import NoPublicConstructor, final, generic_function
 
 if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Awaitable, Callable
     from types import TracebackType
 
-    from typing_extensions import Self
+    from typing_extensions import ParamSpec, Self
+
+    P = ParamSpec("P")
+
+if sys.version_info >= (3, 10):
+    from contextlib import aclosing  # new in Python 3.10
+else:
+
+    class _SupportsAclose(Protocol):
+        def aclose(self) -> Awaitable[object]: ...
+
+    _SupportsAcloseT = TypeVar("_SupportsAcloseT", bound=_SupportsAclose)
+
+    class aclosing(AbstractAsyncContextManager[_SupportsAcloseT, None]):
+        def __init__(self, thing: _SupportsAcloseT) -> None:
+            self._aiter = thing
+
+        async def __aenter__(self) -> _SupportsAcloseT:
+            return self._aiter
+
+        async def __aexit__(self, *exc_info: object) -> None:
+            await self._aiter.aclose()
 
 
 def _open_memory_channel(
@@ -440,3 +467,88 @@ async def aclose(self) -> None:
         See `MemoryReceiveChannel.close`."""
         self.close()
         await trio.lowlevel.checkpoint()
+
+
+def background_with_channel(max_buffer_size: float = 0) -> Callable[
+    [
+        Callable[P, AsyncGenerator[T, None]],
+    ],
+    Callable[P, AbstractAsyncContextManager[trio.MemoryReceiveChannel[T]]],
+]:
+    """Decorate an async generator function to make it cancellation-safe.
+
+    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
+    to the outside `violates structured concurrency <https://xkcd.com/292/>`_
+    with consequences explained in :pep:`789`.  Even then, resource cleanup
+    errors remain common (:pep:`533`) unless you wrap every call in
+    :func:`~contextlib.aclosing`.
+
+    This decorator gives you the best of both worlds: with careful exception
+    handling and a background task we preserve structured concurrency by
+    offering only the safe interface, and you can still write your iterables
+    with the convenience of `yield`.  For example:
+
+        @background_with_channel()
+        async def my_async_iterable(arg, *, kwarg=True):
+            while ...:
+                item = await ...
+                yield item
+
+        async with my_async_iterable(...) as recv_chan:
+            async for item in recv_chan:
+                ...
+
+    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.
+    """
+    # Perhaps a future PEP will adopt `async with for` syntax, like
+    # https://coconut.readthedocs.io/en/master/DOCS.html#async-with-for
+
+    def decorator(
+        fn: Callable[P, AsyncGenerator[T, None]],
+    ) -> Callable[P, AbstractAsyncContextManager[trio.MemoryReceiveChannel[T]]]:
+        @asynccontextmanager
+        @wraps(fn)
+        async def context_manager(
+            *args: P.args, **kwargs: P.kwargs
+        ) -> AsyncGenerator[trio.MemoryReceiveChannel[T], None]:
+            send_chan, recv_chan = trio.open_memory_channel[T](max_buffer_size)
+            async with trio.open_nursery() as nursery:
+                ait = fn(*args, **kwargs)
+                # nursery.start to make sure that we will clean up send_chan & ait
+                await nursery.start(_move_elems_to_channel, ait, send_chan)
+                # async with recv_chan could eat exceptions, so use sync cm
+                with recv_chan:
+                    yield recv_chan
+                # Return promptly, without waiting for `await anext(ait)`
+                nursery.cancel_scope.cancel()
+
+        return context_manager
+
+    async def _move_elems_to_channel(
+        aiterable: AsyncGenerator[T, None],
+        send_chan: trio.MemorySendChannel[T],
+        task_status: trio.TaskStatus,
+    ) -> None:
+        # `async with send_chan` will eat exceptions,
+        # see https://github.com/python-trio/trio/issues/1559
+        with send_chan:
+            async with aclosing(aiterable) as agen:
+                task_status.started()
+                # Outer loop manually advances the aiterable; we can't use async-for because
+                # we're going to use `.asend(err)` to forward errors back to the generator.
+                async for value in agen:
+                    # Get the next value from `agen`; return if exhausted
+                    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
+        # Phew.  Context managers all cleaned up, we're done here.
+
+    return decorator

src/trio/_tests/test_channel.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
 
-from typing import Union
+from typing import TYPE_CHECKING, Union
 
 import pytest
 
 import trio
-from trio import EndOfChannel, open_memory_channel
+from trio import EndOfChannel, background_with_channel, open_memory_channel
 
-from ..testing import assert_checkpoints, wait_all_tasks_blocked
+from ..testing import RaisesGroup, assert_checkpoints, wait_all_tasks_blocked
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator
 
 
 async def test_channel() -> None:
@@ -411,3 +414,107 @@ async def do_send(s: trio.MemorySendChannel[int], v: int) -> None:
             assert await r.receive() == 1
     with pytest.raises(trio.WouldBlock):
         r.receive_nowait()
+
+
+async def test_background_with_channel() -> None:
+    @background_with_channel()
+    async def agen() -> AsyncGenerator[int]:
+        yield 1
+        await trio.sleep_forever()  # simulate deadlock
+
+    async with agen() as recv_chan:
+        async for x in recv_chan:
+            assert x == 1
+            break  # exit, cleanup should be quick
+
+
+async def test_background_with_channel_exhaust() -> None:
+    @background_with_channel()
+    async def agen() -> AsyncGenerator[int]:
+        yield 1
+
+    async with agen() as recv_chan:
+        async for x in recv_chan:
+            assert x == 1
+
+
+async def test_background_with_channel_broken_resource() -> None:
+    @background_with_channel()
+    async def agen() -> AsyncGenerator[int]:
+        yield 1
+        yield 2
+
+    async with agen() as recv_chan:
+        assert await recv_chan.__anext__() == 1
+
+        # close the receiving channel
+        await recv_chan.aclose()
+
+        # trying to get the next element errors
+        with pytest.raises(trio.ClosedResourceError):
+            await recv_chan.__anext__()
+
+        # but we don't get an error on exit of the cm
+
+
+async def test_background_with_channel_cancelled() -> None:
+    with trio.CancelScope() as cs:
+
+        @background_with_channel()
+        async def agen() -> AsyncGenerator[int]:
+            yield 1
+            yield 1
+
+        async with agen():
+            cs.cancel()
+
+
+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()
+    async def agen() -> AsyncGenerator[int]:
+        yield 1
+        raise ValueError("oae")
+
+    with RaisesGroup(ValueError):
+        async with agen() as recv_chan:
+            async for x in recv_chan:
+                assert x == 1
+
+
+async def test_background_with_channel_buffer_size_too_small(
+    autojump_clock: trio.testing.MockClock,
+) -> None:
+    @background_with_channel(0)
+    async def agen() -> AsyncGenerator[int]:
+        yield 1
+        yield 2
+        raise AssertionError(
+            "buffer size 0 means we shouldn't be asked for another value"
+        )  # pragma: no cover
+
+    with trio.move_on_after(5):
+        async with agen() as recv_chan:
+            async for x in recv_chan:
+                assert x == 1
+                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
+        yield 2
+        event.set()
+
+    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__()