A simple task pool for managing and limiting concurrent tasks
| 165 | |
| 166 | |
| 167 | class TaskPool: |
| 168 | """ |
| 169 | A simple task pool for managing and limiting concurrent tasks |
| 170 | """ |
| 171 | |
| 172 | def __init__(self, max_concurrency: int = 10): |
| 173 | self.max_concurrency = max_concurrency |
| 174 | self._semaphore: Optional[asyncio.Semaphore] = None |
| 175 | self.tasks: List[asyncio.Task] = [] |
| 176 | |
| 177 | @property |
| 178 | def semaphore(self) -> asyncio.Semaphore: |
| 179 | """Lazy-initialize the semaphore when first needed""" |
| 180 | if self._semaphore is None: |
| 181 | self._semaphore = asyncio.Semaphore(self.max_concurrency) |
| 182 | return self._semaphore |
| 183 | |
| 184 | async def run(self, coro: Callable, *args: Any, **kwargs: Any) -> Any: |
| 185 | """ |
| 186 | Run a coroutine in the pool |
| 187 | |
| 188 | Args: |
| 189 | coro: Coroutine function to run |
| 190 | *args: Arguments to pass to the coroutine |
| 191 | **kwargs: Keyword arguments to pass to the coroutine |
| 192 | |
| 193 | Returns: |
| 194 | Result of the coroutine |
| 195 | """ |
| 196 | async with self.semaphore: |
| 197 | return await coro(*args, **kwargs) |
| 198 | |
| 199 | def create_task(self, coro: Callable, *args: Any, **kwargs: Any) -> asyncio.Task: |
| 200 | """ |
| 201 | Create and track a task in the pool |
| 202 | |
| 203 | Args: |
| 204 | coro: Coroutine function to run |
| 205 | *args: Arguments to pass to the coroutine |
| 206 | **kwargs: Keyword arguments to pass to the coroutine |
| 207 | |
| 208 | Returns: |
| 209 | Task object |
| 210 | """ |
| 211 | task = asyncio.create_task(self.run(coro, *args, **kwargs)) |
| 212 | self.tasks.append(task) |
| 213 | task.add_done_callback(lambda t: self.tasks.remove(t)) |
| 214 | return task |
| 215 | |
| 216 | async def wait_all(self) -> None: |
| 217 | """Wait for all tasks in the pool to complete""" |
| 218 | if self.tasks: |
| 219 | await asyncio.gather(*self.tasks) |
| 220 | |
| 221 | async def cancel_all(self) -> None: |
| 222 | """Cancel all tasks in the pool""" |
| 223 | for task in self.tasks: |
| 224 | task.cancel() |