fixed issue with not submitting tasks to asyncio.wait. Fixes #35

wolph · May 31, 2022 · d6ade67 · d6ade67
1 parent d768f58
commit d6ade67
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/python_utils/generators.py b/python_utils/generators.py
@@ -6,7 +6,7 @@
 
 
 async def abatcher(
-    generator: types.AsyncGenerator,
+    generator: types.AsyncIterator,
     batch_size: types.Optional[int] = None,
     interval: types.Optional[types.delta_type] = None,
 ):
@@ -33,7 +33,10 @@ async def abatcher(
     while True:
         try:
             done, pending = await asyncio.wait(
-                pending or [generator.__anext__()],
+                pending
+                or [
+                    asyncio.create_task(generator.__anext__()),  # type: ignore
+                ],
                 timeout=interval_s,
                 return_when=asyncio.FIRST_COMPLETED,
             )