| 431 | |
| 432 | |
| 433 | class AsyncGenerator: |
| 434 | def __init__(self, process_func, *args, **kwargs): |
| 435 | self.queue = asyncio.Queue() |
| 436 | self.shutdown_event = threading.Event() |
| 437 | self.iterator_task = None |
| 438 | self.process_func = process_func |
| 439 | self.args = args |
| 440 | self.kwargs = kwargs |
| 441 | |
| 442 | async def producer(self): |
| 443 | # Data generation logic here |
| 444 | for step_result in _generate_tokens( |
| 445 | self.process_func, *self.args, **self.kwargs |
| 446 | ): |
| 447 | await self.queue.put(step_result) |
| 448 | await asyncio.sleep(0.0001) |
| 449 | # asyc sleep otherwise this doesn't yield any result |
| 450 | if self.shutdown_event.is_set(): |
| 451 | break |
| 452 | await self.queue.put(None) |
| 453 | |
| 454 | def __aiter__(self): |
| 455 | self.iterator_task = asyncio.create_task(self.producer()) |
| 456 | return self |
| 457 | |
| 458 | async def __anext__(self): |
| 459 | if self.shutdown_event.is_set(): |
| 460 | raise StopAsyncIteration |
| 461 | |
| 462 | try: |
| 463 | item = await self.queue.get() |
| 464 | if item is None: |
| 465 | self.shutdown_event.set() |
| 466 | raise StopAsyncIteration |
| 467 | return item |
| 468 | except asyncio.CancelledError: |
| 469 | self.shutdown_event.set() |
| 470 | raise StopAsyncIteration |
| 471 | |
| 472 | |
| 473 | def _generate_tokens(process_func, *args, **kwargs): |
no outgoing calls
no test coverage detected