Take n Items Across Chunks Without Reading Ahead

A report reads records that arrive in chunks, and it needs only the first n. Write take_flattened, which returns the first n integers across all the chunks, in order, and starts no chunk after it has enough.

take_flattened([[1, 2], [3, 4], [5, 6]], 4)
# [1, 2, 3, 4], and the third chunk is never started

What to write

def take_flattened(chunks: Iterable[Iterable[int]], n: int) -> list[int]: ...

• Return the first n items across the chunks, in order, as a list. If there are fewer than n, return them all.
• Accept a one-shot iterable of chunks, such as a generator, as well as a list. Ask for each chunk only when you need its items.
• Stop as soon as you have n items. Do not start a later chunk. Do not read past the last item you need, not even to check for empty chunks at the end.
• Use chain.from_iterable to join the chunks and islice to stop after n items, then build a list of only those items. chain(*chunks) starts every chunk before it yields anything, and so does building a full list and slicing it.
• Leave the input values unchanged.

n is 0 or more. The tests never pass a negative n.

What solve reports

The supplied solve(chunks, n) wraps the chunks in a generator, counted(), that adds one to a counter right before it yields each chunk. It passes that generator to take_flattened and returns (items, chunks_started), with the items as a list. The count includes empty chunks started on the way to the next item.

Edge cases

• When n is larger than the number of items, start every chunk and return all the items.
• When n equals the number of items, stop after the last one, without starting the empty chunks after it.
• No chunks at all gives ([], 0).
• With n = 0, start no chunks, even if every chunk is empty.
• With n above 0 and only empty chunks, start every chunk.

The starter returns the right items for finite input, but it reads every chunk before it returns. That is the behavior to fix.