diff --git a/python_utils/generators.py b/python_utils/generators.py index 7f0edf7..ccd8262 100644 --- 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, )