Feature suggestion: strict parameter for zip

[smheidrich]

Python 3.10 introduced the strict parameter for the built-in zip function, which, when set to True, can help spot bugs by raising an exception when the supplied iterables turn out to have different lengths (i.e. when one raises StopIteration after fewer iterations than the others).

I think it would be useful if aiostream's zip function supported this parameter as well.

Extra context: The strict parameter of Python's built-in zip turned out to be useful enough that linters like Ruff or flake8 with the bugbear plugin even have a rule to complain when you neither use it nor explicitly disable it by setting strict=False.

[vxgmichel]

Hi @smheidrich, thanks for the suggestion 👍

I think it would be useful too, feel free to create a PR if you feel like contributing to the project. The implementation is over there:

aiostream/aiostream/stream/combine.py

Lines 49 to 86 in 5e0547a

	@sources_operator
	async def zip(*sources: AsyncIterable[T]) -> AsyncIterator[tuple[T, ...]]:
	"""Combine and forward the elements of several asynchronous sequences.
	Each generated value is a tuple of elements, using the same order as
	their respective sources. The generation continues until the shortest
	sequence is exhausted.
	Note: the different sequences are awaited in parrallel, so that their
	waiting times don't add up.
	"""
	# No sources
	if not sources:
	return
	# One sources
	if len(sources) == 1:
	(source,) = sources
	async with streamcontext(source) as streamer:
	async for item in streamer:
	yield (item,)
	return
	# N sources
	async with AsyncExitStack() as stack:
	# Handle resources
	streamers = [
	await stack.enter_async_context(streamcontext(source)) for source in sources
	]
	# Loop over items
	while True:
	try:
	coros = builtins.map(anext, streamers)
	items = await asyncio.gather(*coros)
	except StopAsyncIteration:
	break
	else:
	yield tuple(items)

It relies on asyncio.gather which exposes a return_exceptions argument. However, I don't think we can use that since we still want to propagate the exception as soon as possible if it is not a StopAsyncIteration. Instead, we probably want to turn StopAsyncIteration into a sentinel with a dedicated function:

async_stop_iteration_sentinel = object()
if strict:
    async def next_item(streamer):
        try:
            return await anext(streamer)
        except StopAsyncIteration:
            return async_stop_iteration_sentinel 
else:
    next_item = anext

[...]
try:
    coros = builtins.map(next_item, streamers)
    items = await asyncio.gather(*coros)
except StopAsyncIteration:
    break
if strict:
    if all(item == async_stop_iteration_sentinel for item in items):
        break
    if any(item == async_stop_iteration_sentinel for item in items):
        raise ValueError
[...]    

Otherwise, I'll do it myself when I find the time :)

[vxgmichel]

Done and released in v0.6.2 🎉

Thanks a lot @smheidrich for your valuable contribution !

[smheidrich]

Great!! Thanks for the reviews & merge! 👍