Batch an Async Stream and Close Its Source

Write an async version of itertools.batched. batched_stream(source, size) reads an async generator and yields lists of size items. It checks size as soon as you call it. It closes source when it stops, whether the source runs out, fails, or the caller closes the stream early.

# source is an async generator that yields 1 to 7
async with aclosing(batched_stream(source, 3)) as batches:
    async for batch in batches:
        print(batch)  # [1, 2, 3], then [4, 5, 6], then [7]

What to write

batched_stream(source, size) is a plain def, so its check runs at the call:

• If size <= 0, raise ValueError(f"size must be positive, got {size}") straight away. Don't read or close source.
• Otherwise return _batched(source, size). Return the generator itself, not a coroutine the caller has to await.

This is the same split as bulk_pages. A check in the body of an async generator would run only at the first anext().

_batched(source, size) is the async generator that does the work:

• Read source in order inside async with aclosing(source). Then the source closes when _batched finishes or raises, and when the caller closes _batched.
• Yield a new list each time you have size items.
• When the source runs out, yield the leftover items as one shorter list, if there are any. An empty source yields nothing.
• Read no further ahead than the batch you are filling.
• If the source raises while you fill a batch, let the error propagate and don't yield the partial batch. The caller keeps the batches it already got.

The yield inside async with aclosing(...) is safe. aclosing sets no timer and cancels nothing, so it is not a cancel scope. Keep timeouts and TaskGroups out of this generator, as the previous article explains.

What the driver reports

Keep _source, _drive, solve and Report as they are. _source yields values in order and logs "pull" for each value it hands out. When the index of the next value equals fail_at, it raises RuntimeError(f"source failed at {index}") instead. A fail_at of None, a negative number or a number past the end never fails. The source's finally logs "closed", so once it has started it logs that even for an empty list.

_drive loops over your stream inside async with aclosing(stream). If the test sets stop_after_batches, the driver stops once it has that many batches. It catches your ValueError and the source's RuntimeError.

• batches: the batches the driver received.
• pulled: how many values the source handed out, counting any in a dropped partial batch.
• source_closed: whether the source logged "closed".
• error: the message of the ValueError or RuntimeError, or "".

Some examples:

• Seven values with size=3 and fail_at=4 give one batch, [1, 2, 3]. The source hands out four values. The fourth goes nowhere, because its batch never filled. The source closes, and error is "source failed at 4".
• With size=2 and stop_after_batches=1, the source hands out two values and closes before solve returns.
• With size=0 there are no batches, pulled is 0, source_closed is False and error is "size must be positive, got 0".
• With a valid size, source_closed is True even for an empty list or a failure at index 0, because _batched starts the source.