如何在Python asyncio中实现等待前N个任务完成？

原代码存在几个关键问题导致无法正常工作：

• 误用了asyncio.await，正确的异步等待API是asyncio.wait
• 返回值不符合需求，没法拿到已完成的任务集合，只返回了剩余任务
• 计数和循环逻辑可以更简洁高效

下面是修正后的实现，能准确返回前N个完成的任务集合，以及剩余未完成的任务：

import asyncio

async def partial_await(task_list, n=None):
    total_tasks = len(task_list)
    # 处理n的边界情况：默认等待全部任务，n<=0或超过总任务数也等待全部
    if n is None or n <= 0 or n > total_tasks:
        n = total_tasks

    done = set()
    pending = set(task_list)

    while len(done) < n and pending:
        # 每次只等第一个完成的任务
        current_done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
        done.update(current_done)

    return done, pending

使用示例 #

比如我们创建5个不同延迟的任务，等待前3个完成：

async def test_task(task_id, delay):
    await asyncio.sleep(delay)
    return f"任务{task_id}完成"

async def main():
    tasks = [
        asyncio.create_task(test_task(1, 3)),
        asyncio.create_task(test_task(2, 1)),
        asyncio.create_task(test_task(3, 2)),
        asyncio.create_task(test_task(4, 0.5)),
        asyncio.create_task(test_task(5, 1.5))
    ]
    
    done, pending = await partial_await(tasks, n=3)
    
    # 打印已完成任务的结果
    print("已完成任务：")
    for task in done:
        print(task.result())
    
    # 可以选择取消剩余未完成的任务
    print("\n取消剩余任务：")
    for task in pending:
        task.cancel()

asyncio.run(main())

逻辑说明 #

1. 初始化done集合存储已完成任务，pending集合存储待处理的任务
2. 循环条件：只要已完成任务数还没到N，且还有未处理任务，就继续等待
3. 每次调用asyncio.wait获取当前第一个完成的任务，更新done和pending
4. 达到N个完成任务时，返回两个集合，方便后续处理已完成结果或取消剩余任务

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