Bound asyncio Concurrency with a Semaphore
Limit concurrent asyncio tasks with a semaphore, keep input order in the results, and report per-task failures without cancelling the batch.
Code
import asyncio
import random
from collections.abc import Awaitable, Callable, Iterable
from dataclasses import dataclass
from typing import Generic, TypeVar
T = TypeVar("T")
R = TypeVar("R")
@dataclass(frozen=True)
class Outcome(Generic[R]):
value: R | None
error: BaseException | None
@property
def ok(self) -> bool:
return self.error is None
async def bounded_map(
items: Iterable[T],
worker: Callable[[T], Awaitable[R]],
*,
limit: int,
) -> list[Outcome[R]]:
if limit < 1:
raise ValueError("limit must be at least 1")
gate = asyncio.Semaphore(limit)
async def guarded(item: T) -> R:
async with gate:
return await worker(item)
outcomes = await asyncio.gather(
*(guarded(item) for item in items),
return_exceptions=True,
)
return [
Outcome(value=None, error=item)
if isinstance(item, BaseException)
else Outcome(value=item, error=None)
for item in outcomes
]
async def fetch(record_id: int) -> dict[str, int]:
await asyncio.sleep(random.uniform(0.01, 0.04))
if record_id == 4:
raise RuntimeError(f"record {record_id} is unavailable")
return {"id": record_id, "status": 200}
async def main() -> None:
outcomes = await bounded_map(range(1, 8), fetch, limit=3)
for outcome in outcomes:
if outcome.ok:
print("ok:", outcome.value)
else:
print("failed:", outcome.error)
if __name__ == "__main__":
asyncio.run(main())
Report a problem with this exampleOpens a pre-filled issue on GitHub.