asynchronous python itertools chain multiple generators
Python's next
built-in function is just a convenient way of invoking the underlying __next__
method on the object. The async equivalent of __next__
is the __anext__
method on the async iterator. There is no anext
global function in the standard library (the aiostream library provides one), but one could easily write it:
async def anext(aiterator):
return await aiterator.__anext__()
But the savings is so small that, in rare situations when this is needed, one may as well invoke __anext__
directly. The async iterator is in turn obtained from an async iterable by calling the __aiter__
(in analogy to __iter__
provided by regular iterables). Async iteration driven manually looks like this:
a_iterator = obj.__aiter__() # regular method
elem1 = await a_iterator.__anext__() # async method
elem2 = await a_iterator.__anext__() # async method
...
__anext__
will raise StopAsyncIteration
when no more elements are available. To loop over async iterators one should use async for
.
Here is a runnable example, based on your code, using both __anext__
and async for
to exhaust the stream set up with aiostream.stream.combine.merge
:
async def main():
a_mix = stream.combine.merge(gen1(), gen2())
async with a_mix.stream() as streamer:
mix_iter = streamer.__aiter__()
print(await mix_iter.__anext__())
print(await mix_iter.__anext__())
print('remaining:')
async for x in mix_iter:
print(x)
asyncio.get_event_loop().run_until_complete(main())
I came across this answer and I looked at the aiostream library. Here is the code I came up with to merge multiple async generators. It does not use any library.
async def merge_generators(gens:Set[AsyncGenerator[Any, None]]) -> AsyncGenerator[Any, None]:
pending = gens.copy()
pending_tasks = { asyncio.ensure_future(g.__anext__()): g for g in pending }
while len(pending_tasks) > 0:
done, _ = await asyncio.wait(pending_tasks.keys(), return_when="FIRST_COMPLETED")
for d in done:
try:
result = d.result()
yield result
dg = pending_tasks[d]
pending_tasks[asyncio.ensure_future(dg.__anext__())] = dg
except StopAsyncIteration as sai:
print("Exception in getting result", sai)
finally:
del pending_tasks[d]
Hope this helps you and let me know if there are any bugs in this.