python 简易协程池

import asyncio
from inspect import iscoroutinefunction
from typing import Any, Awaitable, List
​
​
class CoroutinePool:
    def __init__(self, max_size: int):
        assert max_size > 0, "max_size需大于0"
        self.max_size = max_size
        self._coroutines: List[Awaitable] = []  # 未包装协程
        self._coroutines_with_semaphore: List[Awaitable] = []  # 完成信号量包装协程
​
    async def _semaphore_task(self, semaphore, coroutine) -> Any:
        """
        添加信号量限制
        """
        async with semaphore:
            return await coroutine
​
    def add_coroutine(self, func, *args, **kwargs):
        """
        用户添加协程
        """
        if not iscoroutinefunction(func):
            raise ValueError("只能添加协程函数")
        self._coroutines.append(func(*args, **kwargs))
​
    async def run(self, return_exceptions=False) -> List[Any]:
        """
        执行任务
        return_exceptions :
            bool
            当为 True 时，异常会被返回，而不是被抛出。
            默认值为 False, 当有异常时，会将异常传播到调用方，后续有任务也不继续执行。
        """
        semaphore = asyncio.Semaphore(self.max_size)
        for coroutine in self._coroutines:
            self._coroutines_with_semaphore.append(
                self._semaphore_task(semaphore, coroutine)
            )
        res_list = await asyncio.gather(
            *self._coroutines_with_semaphore, return_exceptions=return_exceptions
        )
        return list(res_list)
​