From 37f74e49758f7be6928ba782c9e875f74819efee Mon Sep 17 00:00:00 2001
From: tongwentao <twtzs@foxmail.com>
Date: Thu, 22 Oct 2026 16:12:01 +0800
Subject: [PATCH] fix: reject nonpositive batch sizes before consuming items

---
 CHANGELOG.md                           |  7 +++++
 _python_utils_tests/test_generators.py | 43 ++++++++++++++++++++++++++
 python_utils/generators.py             | 20 ++++++++++--
 3 files changed, 68 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4a559f0..e474c71 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -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
diff --git a/_python_utils_tests/test_generators.py b/_python_utils_tests/test_generators.py
index 19a9bcc..f0ca8cb 100644
--- a/_python_utils_tests/test_generators.py
+++ b/_python_utils_tests/test_generators.py
@@ -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 = []
diff --git a/python_utils/generators.py b/python_utils/generators.py
index 093cbf3..0794601 100644
--- a/python_utils/generators.py
+++ b/python_utils/generators.py
@@ -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.
 
@@ -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'
@@ -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)