在 Python 中附加到合并的异步生成器

Question

我正在尝试在 Python 3.7 中合并一堆异步生成器，同时仍然在迭代中添加新的异步生成器。我目前正在使用aiostream合并我的生成器：

from asyncio import sleep, run
from aiostream.stream import merge

async def go():
    yield 0
    await sleep(1)
    yield 50
    await sleep(1)
    yield 100

async def main():
    tasks = merge(go(), go(), go())

    async for v in tasks:
        print(v)

if __name__ == '__main__':
    run(main())

但是，一旦循环开始，我需要能够继续添加到正在运行的任务中。就像是。

from asyncio import sleep, run
from aiostream.stream import merge

async def go():
    yield 0
    await sleep(1)
    yield 50
    await sleep(1)
    yield 100

async def main():
    tasks = merge(go(), go(), go())

    async for v in tasks:
        if v == 50:
            tasks.merge(go())
        print(v)

if __name__ == '__main__':
    run(main())

我最接近的是使用aiostream库，但也许这也可以用原生asyncio标准库相当整齐地编写。

Answer 1

这是一个即使使用大量异步迭代器也应该高效工作的实现：

class merge:
    def __init__(self, *iterables):
        self._iterables = list(iterables)
        self._wakeup = asyncio.Event()

    def _add_iters(self, next_futs, on_done):
        for it in self._iterables:
            it = it.__aiter__()
            nfut = asyncio.ensure_future(it.__anext__())
            nfut.add_done_callback(on_done)
            next_futs[nfut] = it
        del self._iterables[:]
        return next_futs

    async def __aiter__(self):
        done = {}
        next_futs = {}
        def on_done(nfut):
            done[nfut] = next_futs.pop(nfut)
            self._wakeup.set()

        self._add_iters(next_futs, on_done)
        try:
            while next_futs:
                await self._wakeup.wait()
                self._wakeup.clear()
                for nfut, it in done.items():
                    try:
                        ret = nfut.result()
                    except StopAsyncIteration:
                        continue
                    self._iterables.append(it)
                    yield ret
                done.clear()
                if self._iterables:
                    self._add_iters(next_futs, on_done)
        finally:
            # if the generator exits with an exception, or if the caller stops
            # iterating, make sure our callbacks are removed
            for nfut in next_futs:
                nfut.remove_done_callback(on_done)

    def append_iter(self, new_iter):
        self._iterables.append(new_iter)
        self._wakeup.set()

示例代码所需的唯一更改是该方法被命名为append_iter，而不是merge.