如何將異步生成器合並到python 3.5+中的vanilla生成器中

Question

我在組合異步生成器和實際運行它們時遇到了麻煩。 這是因為我發現運行它們的唯一方法是通過一個返回可迭代而不是生成器的事件循環。 讓我用一個簡單的例子來說明這一點：

假設我有一個google_search函數，可以通過搜索來搜索谷歌（我沒有故意使用API​​）。 它接收搜索字符串並返回搜索結果的生成器。 當頁面結束時，此生成器不會結束，該功能將繼續到下一頁。 因此google_search函數返回一個可能幾乎無窮無盡的生成器（它在技術上總是會結束，但通常你可以在谷歌搜索獲得數百萬次點擊）

def google_search(search_string):
    # Basically uses requests/aiohttp and beautifulsoup
    # to parse the resulting html and yield search results
    # Assume this function works
    ......

好的，現在我想創建一個允許我迭代多個google_search生成器的函數。 我想要這樣的事情：

def google_searches(*search_strings):
    for results in zip(google_search(query) for query in search_strings):
        yield results

這樣我就可以使用一個簡單的for循環來展開google_searches並獲得我的結果。 上面的代碼運行良好，但對於任何相當多的搜索都非常慢。 代碼正在發送第一次搜索的請求，然后是第二次搜索等等，直到最后，它產生結果。 我想加快速度（很多）。 我的第一個想法是將google_searches更改為異步函數（我使用的是python 3.6.3並且可以使用await / async等）。 然后創建一個異步生成器，但是我只能在另一個異步函數或事件循環中運行它。 使用run_until_complete（loop.gather（...））在事件循環中運行它會返回結果列表而不是普通生成器，這會使目的失敗，因為可能會有太多搜索結果保留在列表中。

如何通過異步執行請求同時使其成為一個香草生成器，使得google_searches功能更快（優選使用異步代碼，但歡迎任何事情）？ 提前致謝！

Answer 1

接受的答案在再次調用生成器之前等待來自EACH異步生成器的一個結果。 如果數據的速度不同，則可能存在問題。 下面的解決方案需要多個異步迭代（生成器或不生成器），並在多個協同程序中同時迭代它們。 每個協程都將結果放在asyncio.Queue ，然后由客戶端代碼迭代：

迭代器代碼：

import asyncio
from async_timeout import timeout

class MergeAsyncIterator:
    def __init__(self, *it, timeout=60, maxsize=0):
        self._it = [self.iter_coro(i) for i in it]
        self.timeout = timeout
        self._futures = []
        self._queue = asyncio.Queue(maxsize=maxsize)

    def __aiter__(self):
        for it in self._it:
            f = asyncio.ensure_future(it)
            self._futures.append(f)
        return self

    async def __anext__(self):
        if all(f.done() for f in self._futures) and self._queue.empty():
            raise StopAsyncIteration
        with timeout(self.timeout):
            try:
                return await self._queue.get()
            except asyncio.CancelledError:
                raise StopAsyncIteration

    def iter_coro(self, it):
        if not hasattr(it, '__aiter__'):
            raise ValueError('Object passed must be an AsyncIterable')
        return self.aiter_to_queue(it)

    async def aiter_to_queue(self, ait):
        async for i in ait:
            await self._queue.put(i)
            await asyncio.sleep(0)

示例客戶端代碼：

import random
import asyncio
from datetime import datetime

async def myaiter(name):
    for i in range(5):
        n = random.randint(0, 3)
        await asyncio.sleep(0.1 + n)
        yield (name, n)
    yield (name, 'DONE')

async def main():
    aiters = [myaiter(i) for i in 'abc']
    async for i in MergeAsyncIterator(*aiters, timeout=3):
        print(datetime.now().strftime('%H:%M:%S.%f'), i)

loop = asyncio.get_event_loop()
loop.run_until_complete(main())

輸出：

14:48:28.638975 ('a', 1)
14:48:29.638822 ('b', 2)
14:48:29.741651 ('b', 0)
14:48:29.742013 ('a', 1)
14:48:30.639588 ('c', 3)
14:48:31.742705 ('c', 1)
14:48:31.847440 ('b', 2)
14:48:31.847828 ('a', 2)
14:48:31.847960 ('c', 0)
14:48:32.950166 ('c', 1)
14:48:33.948791 ('a', 2)
14:48:34.949339 ('b', 3)
14:48:35.055487 ('c', 2)
14:48:35.055928 ('c', 'DONE')
14:48:36.049977 ('a', 2)
14:48:36.050481 ('a', 'DONE')
14:48:37.050415 ('b', 2)
14:48:37.050966 ('b', 'DONE')

PS：上面的代碼使用async_timeout第三方庫。
PS2： aiostream庫與上面的代碼完全相同。

Answer 2

def google_search(search_string):
    # Basically uses requests/aiohttp and beautifulsoup

這是普通的同步發電機。 您可以在其中使用requests ，但是如果您想使用異步aiohttp ，則需要使用async def定義異步生成器 。

什么來迭代多個異步生成器它更有趣。 你不能使用普通的zip因為它適用於普通的迭代，而不是異步迭代。 所以你應該實現自己的（也支持並發迭代）。

我制作了一個我認為可以做你想要的原型：

import asyncio
import aiohttp
import time


# async versions of some builtins:
async def anext(aiterator):
    try:
        return await aiterator.__anext__()
    except StopAsyncIteration as exc:
        raise exc


def aiter(aiterable):
    return aiterable.__aiter__()


async def azip(*iterables):
    iterators = [aiter(it) for it in iterables]
    while iterators:
        results = await asyncio.gather(
            *[anext(it) for it in iterators],
            return_exceptions=True,
        )
        yield tuple(results)


# emulating grabbing:
async def request(url):
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as resp:
            return await resp.text()


async def google_search(search_string):
    for i in range(999):  # big async generator
        url = 'http://httpbin.org/delay/{}'.format(i)  # increase delay to better see concurency
        j = await request(url)
        yield search_string + ' ' + str(i)


async def google_searches(*search_strings):
    async for results in azip(*[google_search(s) for s in search_strings]):
        for result in results:
            yield result


# test it works:
async def main():
    async for result in google_searches('first', 'second', 'third'):
        print(result, int(time.time()))


loop = asyncio.get_event_loop()
try:
    loop.run_until_complete(main())
    loop.run_until_complete(loop.shutdown_asyncgens())
finally:
    loop.close()

輸出：

first 0 1514759561
second 0 1514759561
third 0 1514759561
first 1 1514759562
second 1 1514759562
third 1 1514759562
first 2 1514759564
second 2 1514759564
third 2 1514759564
first 3 1514759567
second 3 1514759567
third 3 1514759567

時間顯示不同的搜索同時運行。

Answer 3

我只是要粘貼我剛才編寫的解決方案，因為我總是在這個問題中結束，只是為了記住我之前已經解決了這個問題。

async def iterator_merge(iterators: typing.Dict[typing.AsyncIterator, typing.Optional[asyncio.Future]]):
while iterators:
    for iterator, value in list(iterators.items()):
        if not value:
            iterators[iterator] = asyncio.ensure_future(iterator.__anext__())

    tasks, _ = await asyncio.wait(iterators.values(), return_when=asyncio.FIRST_COMPLETED)
    for task in tasks:
        # We send the result up
        try:
            res = task.result()
            yield res
        except StopAsyncIteration:
            # We remove the task from the list
            for it, old_next in list(iterators.items()):
                if task is old_next:
                    logger.debug(f'Iterator {it} finished consuming')
                    iterators.pop(it)
        else:
            # We remove the task from the key
            for it, old_next in list(iterators.items()):
                if task is old_next:
                    iterators[it] = None

它有打字注釋，但我認為這是一個很好的解決方案。 它意味着使用異步生成器作為鍵調用，如果有任何等待，則需要將來調用。

iterators = {
    k8s_stream_pod_log(name=name): None,
    k8s_stream_pod_events(name=name): None,
}

您可以在github.com/txomon/abot中找到它的使用方法。