diff --git a/CHANGELOG.md b/CHANGELOG.md index 73b9cf2..769cbf1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +## Unreleased + +- Flush a nonempty `abatcher` batch when its interval expires instead of + waiting a full interval again for each item. An arrival exactly at the + deadline now flushes the batch as well. + ## 4.1.2 - 2026-10-13 - Count the time that passed across a daylight saving change in `timesince` diff --git a/_python_utils_tests/test_generators.py b/_python_utils_tests/test_generators.py index 501db0f..9fdd0a8 100644 --- a/_python_utils_tests/test_generators.py +++ b/_python_utils_tests/test_generators.py @@ -1,6 +1,7 @@ """Tests for the batching helpers in ``python_utils.generators``.""" import asyncio +import datetime import gc import sys from types import SimpleNamespace @@ -23,14 +24,22 @@ async def test_abatcher() -> None: @pytest.mark.asyncio @pytest.mark.parametrize( - 'arrival_times', + ('arrival_times', 'expected'), [ - (0, 8, 16, 24, 32, 40, 48, 56, 64, 72), - (0, 10, 11, 21, 22, 32, 33, 43, 44, 54), + ( + (0, 8, 16, 24, 32, 40, 48, 56, 64, 72), + [[0, 1, 2], [3, 4], [5, 6], [7, 8], [9]], + ), + ( + (0, 10, 11, 21, 22, 32, 33, 43, 44, 54), + [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]], + ), ], ) async def test_abatcher_timed( - monkeypatch: pytest.MonkeyPatch, arrival_times: tuple[int, ...] + monkeypatch: pytest.MonkeyPatch, + arrival_times: tuple[int, ...], + expected: list[list[int]], ) -> None: """Group async items into batches by time interval.""" now: float = 0.0 @@ -52,7 +61,7 @@ async def generator() -> types.AsyncIterator[int]: async for batch in python_utils.abatcher(generator(), interval=10): batches.append(batch) - assert batches == [[0, 1, 2], [3, 4], [5, 6], [7, 8], [9]] + assert batches == expected assert len(batches) == 5 @@ -89,6 +98,61 @@ async def generator() -> types.AsyncIterator[int]: await batcher.__anext__() +@pytest.mark.asyncio +@pytest.mark.parametrize('batch_size', [None, 3]) +@pytest.mark.parametrize('interval', [10, datetime.timedelta(seconds=10)]) +@pytest.mark.parametrize('arrival', [8.0, 10.0]) +async def test_abatcher_partial_batch_keeps_deadline( + monkeypatch: pytest.MonkeyPatch, + batch_size: int | None, + interval: int | datetime.timedelta, + arrival: float, +) -> None: + """Flush a partial batch at its deadline, including an exact arrival.""" + now: float = 0.0 + monkeypatch.setattr( + python_utils.generators, + 'time', + SimpleNamespace(perf_counter=lambda: now), + ) + original_wait = asyncio.wait + + async def wait( + pending: set[asyncio.Future[int]], + *, + timeout: float, + return_when: str, + ) -> tuple[set[asyncio.Future[int]], set[asyncio.Future[int]]]: + """Run ready items, or let the requested timeout elapse.""" + nonlocal now + done, remaining = await original_wait( + pending, timeout=0, return_when=return_when + ) + if not done: + now += timeout + return done, remaining + + monkeypatch.setattr(asyncio, 'wait', wait) + + async def generator() -> types.AsyncGenerator[int, None]: + """Yield twice on the fake clock, then leave one item pending.""" + nonlocal now + yield 0 + now = arrival + yield 1 + await asyncio.Event().wait() + + batcher = python_utils.abatcher( + generator(), batch_size=batch_size, interval=interval + ) + try: + first = await batcher.__anext__() + assert first == [0, 1] + assert now == 10.0 + finally: + await batcher.aclose() + + async def blocked_generator( started: asyncio.Event, closed: types.List[str] ) -> types.AsyncIterator[int]: diff --git a/python_utils/generators.py b/python_utils/generators.py index e6571bb..30fc1a7 100644 --- a/python_utils/generators.py +++ b/python_utils/generators.py @@ -72,6 +72,9 @@ async def abatcher( request for it. A source that is an async generator ends at that point. A source with nothing on its way stays usable. + Items arriving within an interval share its deadline. Waiting for the + next item does not restart the interval of a nonempty batch. + Args: generator: The async generator or iterator to batch. batch_size (typing.Optional[int], optional): The number of items per @@ -110,9 +113,13 @@ async def abatcher( # `create_task` insists on a coroutine. pending = {asyncio.ensure_future(generator.__anext__())} + timeout = interval_s + if interval and batch: + timeout = max(0.0, next_yield - time.perf_counter()) + done, pending = await asyncio.wait( pending, - timeout=interval_s, + timeout=timeout, return_when=asyncio.FIRST_COMPLETED, ) @@ -132,7 +139,7 @@ async def abatcher( # next item is flushed on its own once the old one ran out. next_yield = time.perf_counter() + interval_s - if interval and batch and time.perf_counter() > next_yield: + if interval and batch and time.perf_counter() >= next_yield: yield batch batch = [] # Always set the next yield time to the current time. If the