Python中实现无同步异步迭代器zip函数async_zip_stream的疑问

要实现能无同步遍历多个异步迭代器、缺省值填充的async_zip_stream函数，核心思路是为每个异步迭代器单独维护取值任务，监听任意迭代器的产出事件，一旦有值产出就立即生成对应元组。以下是具体实现：

import asyncio

async def async_zip_stream(*args, default_value=0):
    # 将传入的异步生成器转为异步迭代器
    iters = [aiter(arg) for arg in args]
    # 维护每个迭代器的任务、完成状态及索引
    tasks = []
    for idx, it in enumerate(iters):
        # 启动第一个取值任务
        task = asyncio.create_task(anext(it))
        tasks.append((task, False, idx))
    
    while tasks:
        # 等待任意一个迭代器完成取值或耗尽
        done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
        for completed_task in done:
            # 定位完成的任务项
            task_info_idx = next(i for i, (t, _, idx) in enumerate(tasks) if t is completed_task)
            task, _, idx = tasks.pop(task_info_idx)
            
            try:
                value = completed_task.result()
                # 生成元组：仅当前产出值的位置用实际值，其余用缺省值
                yield tuple(
                    value if i == idx else default_value
                    for i in range(len(args))
                )
                # 为当前迭代器启动下一个取值任务
                next_task = asyncio.create_task(anext(iters[idx]))
                tasks.append((next_task, False, idx))
            except StopAsyncIteration:
                # 迭代器已耗尽，不再启动新任务
                pass

# 测试用例
async def iter1():
    for i in range(5):
        await asyncio.sleep(0.1)
        yield 1

async def iter2():
    for i in range(10):
        await asyncio.sleep(1)
        yield 2

async def main():
    async for item in async_zip_stream(iter1(), iter2()):
        print(", ".join(map(str, item)))

if __name__ == "__main__":
    asyncio.run(main())

代码说明 #

1. 初始化任务：将每个输入的异步生成器转为异步迭代器，并为每个迭代器启动第一个anext()任务，异步获取第一个值。
2. 监听产出事件：通过asyncio.wait的FIRST_COMPLETED模式，等待任意一个迭代器完成取值或耗尽，实现无同步响应。
3. 处理产出值：当某个迭代器产出值时，生成对应元组（仅该位置用实际值，其余用缺省值）并yield；同时为该迭代器启动下一个取值任务，持续获取后续值。
4. 处理迭代器耗尽：若迭代器抛出StopAsyncIteration，说明已无更多值，不再为其启动新任务，后续元组中该位置始终使用缺省值。

运行测试代码后，会先每0.1秒输出一次1, 0（共5次），之后每1秒输出一次0, 2（共10次），完全符合需求。

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