是否存在与异步迭代器一起使用的itertools.chain(* iterables)的等价物?一个关键要求是尽快从异步迭代器获取可用数据(即没有天真的链接)。
更新:请注意,重复问题的一个关键区别是,下面的答案可让您识别触发异步生成器。
答案 0 :(得分:2)
以下代码解决了这个问题:
import asyncio
class InternalStopAsyncIteration(Exception):
"""A special stop exception that also returns the finished generator's key."""
def __init__(self, key):
self.key = key
async def anext(key, gen):
try:
return key, await gen.__anext__()
except StopAsyncIteration:
raise InternalStopAsyncIteration(key)
async def combine_async_generators(**gens):
pending = {anext(key, gen) for key, gen in gens.items()}
while pending:
done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
for i in done:
if isinstance(i.exception(), InternalStopAsyncIteration):
gens.pop(i.exception().key)
else:
key, val = i.result()
pending.add(anext(key, gens[key]))
yield key, val
# The following will print:
# a 0.5
# b 1
# a 0.5
# a 0.5
# b 1
# b 1
async def gen(x):
"""An async generator that sleeps a bit, then yields the given value."""
for i in range(3):
await asyncio.sleep(x)
yield x
async def run():
async for k, v in combine_async_generators(a=gen(0.5), b=gen(1)):
print(k, v)
asyncio.get_event_loop().run_until_complete(run())