A simple task pool for managing and limiting concurrent asyncio tasks.
| 196 | |
| 197 | |
| 198 | class TaskPool: |
| 199 | """ |
| 200 | A simple task pool for managing and limiting concurrent asyncio tasks. |
| 201 | """ |
| 202 | |
| 203 | def __init__(self, max_concurrency: int = 10): |
| 204 | """Initializes the task pool with a concurrency limit.""" |
| 205 | self.max_concurrency = max_concurrency |
| 206 | # The semaphore is initialized lazily to ensure it's created within a running event loop |
| 207 | self._semaphore: Optional[asyncio.Semaphore] = None |
| 208 | # A list to keep track of currently running tasks |
| 209 | self.tasks: List[asyncio.Task] = [] |
| 210 | |
| 211 | @property |
| 212 | def semaphore(self) -> asyncio.Semaphore: |
| 213 | """Lazy-initialize the semaphore when first needed.""" |
| 214 | if self._semaphore is None: |
| 215 | self._semaphore = asyncio.Semaphore(self.max_concurrency) |
| 216 | return self._semaphore |
| 217 | |
| 218 | async def run(self, coro: Callable, *args: Any, **kwargs: Any) -> Any: |
| 219 | """ |
| 220 | Run a coroutine within the pool's concurrency limit. |
| 221 | |
| 222 | Args: |
| 223 | coro: The coroutine function to run. |
| 224 | *args: Positional arguments for the coroutine. |
| 225 | **kwargs: Keyword arguments for the coroutine. |
| 226 | |
| 227 | Returns: |
| 228 | The result of the coroutine. |
| 229 | """ |
| 230 | # Acquire the semaphore before running the coroutine |
| 231 | async with self.semaphore: |
| 232 | # Await the coroutine |
| 233 | return await coro(*args, **kwargs) |
| 234 | |
| 235 | def create_task(self, coro: Callable, *args: Any, **kwargs: Any) -> asyncio.Task: |
| 236 | """ |
| 237 | Create and track a task that will run in the pool. |
| 238 | |
| 239 | Args: |
| 240 | coro: The coroutine function to create a task for. |
| 241 | *args: Positional arguments for the coroutine. |
| 242 | **kwargs: Keyword arguments for the coroutine. |
| 243 | |
| 244 | Returns: |
| 245 | The created asyncio.Task object. |
| 246 | """ |
| 247 | # Create a task that runs the coroutine through the pool's 'run' method |
| 248 | task = asyncio.create_task(self.run(coro, *args, **kwargs)) |
| 249 | # Add the task to the tracking list |
| 250 | self.tasks.append(task) |
| 251 | # Add a callback to automatically remove the task from the list when it's done |
| 252 | task.add_done_callback(lambda t: self.tasks.remove(t)) |
| 253 | return task |
| 254 | |
| 255 | async def wait_all(self) -> None: |