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 协程的代码需要额外判断。