Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,5 +1,12 @@
# Changelog

## Unreleased

- Reject zero and negative batch sizes with `ValueError` in `batcher` and
`abatcher` before consuming the source. These sizes never reached the batch
limit and could buffer an unbounded source indefinitely. `abatcher` still
accepts `None` for interval-only batching.

## 4.1.4 - 2026-10-21

- Build from the source distribution with `uv_build` 0.13 and later. The
Expand Down
43 changes: 43 additions & 0 deletions _python_utils_tests/test_generators.py
Original file line number Diff line number Diff line change
Expand Up @@ -364,6 +364,49 @@ async def test_abatcher_with_future_returning_iterator() -> None:
assert batches == [[0, 1], [2, 3], [4]]


@pytest.mark.parametrize('batch_size', [0, -1])
@pytest.mark.parametrize('count', [0, 3])
def test_batcher_rejects_nonpositive_size(batch_size: int, count: int) -> None:
"""Reject invalid sizes without consuming even a finite source."""
consumed: list[int] = []

def source() -> types.Generator[int, None, None]:
for item in range(count):
consumed.append(item)
yield item

with pytest.raises(
ValueError, match='batch_size must be greater than zero'
):
list(python_utils.batcher(source(), batch_size))

assert consumed == []


@pytest.mark.asyncio
@pytest.mark.parametrize('batch_size', [0, -1])
@pytest.mark.parametrize('count', [0, 3])
@pytest.mark.parametrize('interval', [None, 1])
async def test_abatcher_rejects_nonpositive_size(
batch_size: int, count: int, interval: int | None
) -> None:
"""Reject invalid sizes before scheduling a request from the source."""
consumed: list[int] = []

async def source() -> types.AsyncGenerator[int, None]:
for item in range(count):
consumed.append(item)
yield item

with pytest.raises(
ValueError, match='batch_size must be greater than zero'
):
async for _ in python_utils.abatcher(source(), batch_size, interval):
pass

assert consumed == []


def test_batcher() -> None:
"""Split an iterable into fixed-size batches."""
batch = []
Expand Down
20 changes: 18 additions & 2 deletions python_utils/generators.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,12 @@
_T = typing.TypeVar('_T')


def _validate_batch_size(batch_size: int | None) -> None:
"""Require a positive batch size, or None for interval-only batching."""
if batch_size is not None and batch_size < 1:
raise ValueError('batch_size must be greater than zero')


async def _cancel(pending: set[asyncio.Future[_T]]) -> None:
"""Cancel the items that are still on their way and wait for them.

Expand Down Expand Up @@ -78,14 +84,19 @@ async def abatcher(
Args:
generator: The async generator or iterator to batch.
batch_size (typing.Optional[int], optional): The number of items per
batch. Defaults to None.
batch. Must be greater than zero when specified. Defaults to None.
interval (typing.Optional[_aliases.delta_type], optional): The time
interval to wait before yielding a batch. Defaults to None.

Yields:
collections.abc.AsyncGenerator[list[_T], None]: A generator that yields
batches of items.

Raises:
ValueError: If batch_size is specified but is not greater than zero.
"""
_validate_batch_size(batch_size)

batch: list[_T] = []

assert batch_size or interval, 'Must specify either batch_size or interval'
Expand Down Expand Up @@ -170,12 +181,17 @@ def batcher(
Args:
iterable (collections.abc.Iterable[_T]): The iterable to batch.
batch_size (int, optional): The number of items per batch. Defaults
to 10.
to 10. Must be greater than zero.

Yields:
collections.abc.Generator[list[_T], None, None]: A generator that
yields batches of items.

Raises:
ValueError: If batch_size is not greater than zero.
"""
_validate_batch_size(batch_size)

batch: list[_T] = []
for item in iterable:
batch.append(item)
Expand Down