【Python】用装饰器限制协程并发量

from functools import wraps
from random import randint
import asyncio
from typing import Coroutine, List
from asyncio import Task

import nest_asyncio

nest_asyncio.apply()


def limit(max_coro):
    semphore = asyncio.Semaphore(max_coro)

    def _wrapper(func: Coroutine):
        @wraps(func)
        async def _worker(*args, **kwargs):
            async with semphore:
                return await func(*args, **kwargs)

        return _worker

    return _wrapper


# 并发量为 1
@limit(1)
async def request(idx):
    delay = randint(1, 1)
    print(idx, 'now requesting', 'delay', delay)
    await asyncio.sleep(delay)
    print(idx, 'delay', delay, 'done')
    return idx, delay

async def main():
    tasks: List[Task] = []

    for _ in range(6):
        task = asyncio.create_task(request(_))
        tasks.append(task)

    for task in asyncio.as_completed(tasks):
        _idx, delay = await task

        # 这里 running 的协程也会被取消
        # 可能需要用下标等方式标记当前返回的协程
        if delay > 2:
            # print(task.)
            for _task in tasks:
                print(_task.done())
                _task.cancel()
                await asyncio.sleep(0)
                # _task.
                # print('cancelled', '*' * 10)
                # print(_task)

            break

    await asyncio.wait(tasks)

原理很简单，就是装饰器 + 闭包。关于取消 running 协程的代码需要额外判断。