test strict batched mode

Author: maxfischer2781

asyncstdlib/itertools.py

@@ -140,7 +140,7 @@ async def batched(
                 for _ in range(n):
                     batch.append(await anext(item_iter))
             except StopAsyncIteration:
-                if strict and len(batch) < n:
+                if strict and batch and len(batch) < n:
                     raise ValueError("batched(): incomplete batch") from None
             if batch:
                 yield tuple(batch)

unittests/test_itertools.py

@@ -79,6 +79,19 @@ async def test_batched_invalid(length):
         await a.list(a.batched(range(10), length))
 
 
+@sync
+@pytest.mark.parametrize("values", ([1, 2, 3, 4], [1, 2, 3, 4, 5], [1]))
+async def test_batched_strict(values: "list[int]"):
+    for n in range(1, len(values) + 1):
+        batches = a.batched(values, n, strict=True)
+        if len(values) % n == 0:
+            assert values == list(await a.reduce(lambda a, b: a + b, batches))
+        else:
+            assert await a.anext(batches)
+            with pytest.raises(ValueError):
+                await a.list(batches)
+
+
 @sync
 async def test_cycle():
     async for _ in a.cycle([]):