如何将异步可迭代对象转换为非列表可迭代对象并规避语法错误

原代码（运行正常） #

import asyncio
from collections.abc import AsyncIterable, Iterable


# 实际场景中逻辑更复杂
async def create_asynchronous_iterable() -> AsyncIterable[int]:
    for k in range(3):
        yield k


# 实际场景中逻辑更复杂
def do_something(it: Iterable[int]) -> None:
    for k in it:
        print(k, end=" ")
    print()


async def main():
    async_it = create_asynchronous_iterable()
    it = [k async for k in async_it]
    do_something(it)


if __name__ == "__main__":
    asyncio.run(main())  # 输出 "0 1 2 "

问题描述 #

当前代码会把异步迭代器的所有元素一次性加载到列表中再传入do_something，现在需要将it转换为无需预加载所有元素到内存的流式可迭代对象（比如生成器），避免大数量场景下的内存占用问题。

错误尝试及报错 #

尝试在同步生成器中直接使用async for，触发语法错误：

async def main():
    async_it = create_asynchronous_iterable()

    def create_iterable() -> Iterable[int]:
        async for k in async_it:
            yield k
    
    it = create_iterable()
    do_something(it)

报错信息：

SyntaxError: 'async for' outside async function

解决方案 #

核心思路是借助当前运行的事件循环，在同步生成器中逐个异步获取元素并返回，实现流式迭代。修改后的main函数如下：

async def main():
    async_it = create_asynchronous_iterable()
    loop = asyncio.get_running_loop()

    def create_iterable() -> Iterable[int]:
        # 把异步迭代器包装为异步生成器
        async def async_gen():
            async for item in async_it:
                yield item
        
        agen = async_gen()
        while True:
            try:
                # 通过事件循环阻塞获取下一个异步元素
                item = loop.run_until_complete(agen.__anext__())
                yield item
            except StopAsyncIteration:
                break

    it = create_iterable()
    do_something(it)

原理说明 #

1. async for语法仅允许在异步函数/异步生成器中使用，因此先将原异步迭代器包装为异步生成器async_gen。
2. 在同步生成器create_iterable中，通过当前事件循环的run_until_complete()方法，每次阻塞等待获取异步生成器的下一个元素，再通过同步生成器yield出去。
3. 这种方式不会一次性加载所有元素到内存，而是在do_something迭代时才逐个拉取异步元素，实现流式处理。

内容的提问来源于stack exchange，提问作者Sebastian Thomas