Fixed bug with batcher skipping items in the case of timeouts thanks to @jorenham

Author: wolph

_python_utils_tests/test_generators.py

@@ -1,3 +1,5 @@
+import asyncio
+
 import pytest
 
 import python_utils
@@ -16,12 +18,42 @@ async def test_abatcher():
 async def test_abatcher_timed():
     batches = []
     async for batch in python_utils.abatcher(
-        python_utils.acount(stop=10, delay=0.08), interval=0.2
+        python_utils.acount(stop=10, delay=0.08), interval=0.1
     ):
         batches.append(batch)
 
-    assert len(batches) == 3
-    assert sum(len(batch) for batch in batches) == 10
+    assert batches == [[0, 1, 2], [3, 4], [5, 6], [7, 8], [9]]
+    assert len(batches) == 5
+
+
+@pytest.mark.asyncio
+async def test_abatcher_timed_with_timeout():
+    async def generator():
+        # Test if the timeout is respected
+        yield 0
+        yield 1
+        await asyncio.sleep(0.11)
+
+        # Test if the timeout is respected
+        yield 2
+        yield 3
+        await asyncio.sleep(0.11)
+
+        # Test if exceptions are handled correctly
+        await asyncio.wait_for(asyncio.sleep(1), timeout=0.05)
+
+        # Test if StopAsyncIteration is handled correctly
+        yield 4
+
+    batcher = python_utils.abatcher(generator(), interval=0.1)
+    assert await batcher.__anext__() == [0, 1]
+    assert await batcher.__anext__() == [2, 3]
+
+    with pytest.raises(asyncio.TimeoutError):
+        await batcher.__anext__()
+
+    with pytest.raises(StopAsyncIteration):
+        await batcher.__anext__()
 
 
 def test_batcher():

python_utils/generators.py

@@ -18,27 +18,35 @@ async def abatcher(
 
     assert batch_size or interval, 'Must specify either batch_size or interval'
 
+    # If interval is specified, use it to determine when to yield the batch
+    # Alternatively set a really long timeout to keep the code simpler
     if interval:
         interval_s = python_utils.delta_to_seconds(interval)
-        next_yield = time.perf_counter() + interval_s
     else:
-        interval_s = 0
-        next_yield = 0
+        # Set the timeout to 10 years
+        interval_s = 60 * 60 * 24 * 365 * 10.0
+
+    next_yield = time.perf_counter() + interval_s
+
+    pending: types.Set = set()
 
     while True:
         try:
-            if interval_s:
-                item = await asyncio.wait_for(
-                    generator.__anext__(), interval_s
-                )
-            else:
-                item = await generator.__anext__()
-        except (StopAsyncIteration, asyncio.TimeoutError):
+            done, pending = await asyncio.wait(
+                pending or [generator.__anext__()],
+                timeout=interval_s,
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+
+            if done:
+                for result in done:
+                    batch.append(result.result())
+
+        except StopAsyncIteration:
             if batch:
                 yield batch
+
             break
-        else:
-            batch.append(item)
 
         if batch_size is not None and len(batch) == batch_size:
             yield batch