Batched async stream: never yield inside a cancel scope

PEP 789 documents the failure precisely. Yielding out of a frame that has an open TaskGroup or asyncio.timeout scope suspends that frame while leaving the scope open. The scope’s cancellation is then delivered to whichever task happens to be driving the generator next. The symptom is a CancelledError or an ExceptionGroup surfacing in a task that has nothing to do with the code that opened the scope, which is roughly the worst debugging experience asyncio offers.

PEP 789 proposes sys.prevent_yields() to make this a runtime error. It is Draft and has not landed, so today the mitigation is a review rule, not an API. Learn the rule:

Never yield while a TaskGroup or asyncio.timeout scope is open in the same frame.

The shape that lures people in is @asynccontextmanager wrapped around a TaskGroup — it looks like clean resource management and it is a cancellation bug.

aclosing is not a cancel scope. async with contextlib.aclosing(x) is an ordinary async context manager whose __aexit__ awaits x.aclose(). Yielding across it is fine. Learning to tell “ordinary async CM” from “cancel scope” is the actual skill; the review rule is the shortcut.

What to write

def batched_stream[T](source: AsyncGenerator[T, None], size: int) -> AsyncGenerator[list[T], None]
async def _batched[T](source: AsyncGenerator[T, None], size: int) -> AsyncGenerator[list[T], None]

batched_stream is a plain def. It validates size and returns the async generator produced by _batched. _batched is the async def ... yield worker.

Why split it? Because adding a single yield changes the function’s type. An async def with no yield and a declared return of AsyncIterator[T] is a coroutine returning an async iterator — you await it, then iterate what comes back. Add a yield and it becomes an async generator function: calling it returns the iterator directly and runs none of the body. That means an async generator cannot raise at call time, and size <= 0 has to be rejected before the caller starts iterating. Hence the wrapper. mypy will happily tell you which of the two you have written if you ask it to reveal_type.

The contract:

• size <= 0 raises ValueError(f"size must be positive, got {size}") immediately, before source is touched at all.
• Yield lists of exactly size items, then the short remainder if any.
• source is closed on every exit path: normal exhaustion, an early consumer break, and an exception thrown from inside source. Wrap the iteration in async with aclosing(source).
• No TaskGroup, no asyncio.timeout, held open across the yield.

The driver, the instrumented source, and the report are provided.

What the report proves

• batches — exact batching and the short remainder.
• pulled — how many values the source actually produced.
• source_closed — whether the source’s finally ran. It must be True after a normal end, after an early break, and after a mid-stream failure — and False when size <= 0, because a correct implementation never touched the source at all.
• error — the message of whatever escaped, or "".

The early-break case is the one that separates working code from correct code. The consumer breaks after the first batch; aclosing then throws GeneratorExit into _batched at its yield, whose own aclosing(source) __aexit__ closes the source — all before the async with block in the driver has finished unwinding. source_closed being observable synchronously is the whole point. Without aclosing you are relying on the garbage collector and the loop’s shutdown_asyncgens, which is to say on nothing.