utils/async_utils.py

"""Async batching utilities for bounded concurrency."""

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Sequence


async def batched_gather[T](
    awaitables: Sequence[Awaitable[T]],
    batch_size: int,
) -> list[T]:
    """Run awaitables in batches, returning all results in original order.

    Prevents unbounded asyncio.gather from overwhelming SQLite or network resources.
    Each batch is gathered concurrently; batches run sequentially.

    Args:
        awaitables: Sequence of awaitables to run.
        batch_size: Maximum number of awaitables per batch.

    Returns:
        List of results in the same order as input awaitables.
    """
    total = len(awaitables)
    results: list[T] = []
    batch_start = 0
    # nosemgrep: python-sequential-await-in-loop
    while batch_start < total:
        batch_end = min(batch_start + batch_size, total)
        chunk = awaitables[batch_start:batch_end]
        batch_results = await asyncio.gather(*chunk)
        results.extend(batch_results)
        batch_start = batch_end
    return results