如何并行链接多个异步生成器/迭代器?

时间:2018-03-20 10:00:30

标签: python async-await python-asyncio

是否存在与异步迭代器一起使用的itertools.chain(* iterables)的等价物?一个关键要求是尽快从异步迭代器获取可用数据(即没有天真的链接)。

更新:请注意,重复问题的一个关键区别是,下面的答案可让您识别触发异步生成器。

1 个答案:

答案 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())