如何在AnyIO中实现asyncio as_completed的功能？

核心思路 #

AnyIO没有直接提供asyncio.as_completed的等价API，但可以通过循环调用anyio.wait并设置return_when=FIRST_COMPLETED来实现相同的「任务完成即处理」逻辑：每次调用会返回已完成的任务和未完成的任务，处理完已完成任务的结果后，继续等待剩余任务，直到所有任务执行完毕。

修改后的完整代码 #

import itertools
import anyio
from anyio import wait, FIRST_COMPLETED


async def handle(event: str):
    await handle_event(event)


async def handle_event(event: str):
    if event == "event":
        tasks = [anyio.create_task(slow_2()), anyio.create_task(slow_5())]
    else:
        tasks = [anyio.create_task(slow_1()), anyio.create_task(slow_1())]
    
    while tasks:
        # 等待第一个完成的任务
        done, pending = await wait(tasks, return_when=FIRST_COMPLETED)
        tasks = list(pending)  # 更新剩余任务列表
        
        for task in done:
            result = await task
            new_events = []
            if result == "event":
                new_events.append("sub_event")  # 修正原代码事件格式问题，可根据业务需求调整
            
            if new_events:
                async with anyio.create_task_group() as tg:
                    for sub_event in new_events:
                        tg.start_soon(handle_event, sub_event)


async def spin(msg: str) -> None:
    for char in itertools.cycle(r"\|/-/"):
        status = f"\r{char} {msg}"
        print(status, flush=True, end="")
        try:
            await anyio.sleep(0.1)
        except Exception:
            break
    blanks = " " * len(status)
    print(f"\r{blanks}\r", end="")


async def slow_1():
    await anyio.sleep(1)
    print("slow_1")


async def slow_2():
    await anyio.sleep(2)
    print("slow_2")
    return "event"


async def slow_5():
    await anyio.sleep(5)
    print("slow_5")


async def supervisor():
    async with anyio.create_task_group() as tg:
        with anyio.CancelScope(shield=True):
            tg.start_soon(spin, "thinking!")
            await handle("event")
            tg.cancel_scope.cancel()


if __name__ == "__main__":
    anyio.run(supervisor)

关键修改点说明 #

1. 替换asyncio.as_completed逻辑：
   • 先通过anyio.create_task将协程包装成可跟踪的任务对象
   • 循环调用anyio.wait获取首个完成的任务，处理结果后更新剩余任务列表，直到所有任务处理完毕
2. 任务管理适配AnyIO：AnyIO的任务跟踪需要显式创建Task对象，不能直接传入协程列表给wait方法
3. 事件格式修正：原代码中new_events.append(["", ""])会导致事件处理器接收错误类型的参数，这里调整为符合函数参数要求的字符串格式，可根据实际业务修改

备选实现方式（队列收集结果） #

如果需要更灵活的结果处理逻辑，也可以用队列收集任务输出，实现与as_completed一致的顺序处理：

async def handle_event(event: str):
    if event == "event":
        coros = [slow_2(), slow_5()]
    else:
        coros = [slow_1(), slow_1()]
    
    result_queue = anyio.Queue()
    
    async with anyio.create_task_group() as tg:
        # 启动所有任务，完成后将结果写入队列
        for coro in coros:
            async def task_wrapper(c):
                res = await c
                await result_queue.put(res)
            tg.start_soon(task_wrapper, coro)
        
        # 按任务完成顺序从队列取结果处理
        for _ in coros:
            result = await result_queue.get()
            new_events = []
            if result == "event":
                new_events.append("sub_event")
            
            if new_events:
                async with anyio.create_task_group() as sub_tg:
                    for sub_event in new_events:
                        sub_tg.start_soon(handle_event, sub_event)

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