A stream of RequestOutputs for a request that can be iterated over asynchronously.
| 40 | |
| 41 | |
| 42 | class AsyncStream: |
| 43 | """A stream of RequestOutputs for a request that can be |
| 44 | iterated over asynchronously.""" |
| 45 | |
| 46 | def __init__(self, request_id: str) -> None: |
| 47 | self.request_id = request_id |
| 48 | self._queue = asyncio.Queue() |
| 49 | self._finished = False |
| 50 | |
| 51 | def put(self, item: RequestOutput) -> None: |
| 52 | if self._finished: |
| 53 | return |
| 54 | self._queue.put_nowait(item) |
| 55 | |
| 56 | def finish(self) -> None: |
| 57 | self._queue.put_nowait(StopIteration) |
| 58 | self._finished = True |
| 59 | |
| 60 | @property |
| 61 | def finished(self) -> bool: |
| 62 | return self._finished |
| 63 | |
| 64 | def __aiter__(self): |
| 65 | return self |
| 66 | |
| 67 | async def __anext__(self) -> RequestOutput: |
| 68 | result = await self._queue.get() |
| 69 | if result is StopIteration: |
| 70 | raise StopAsyncIteration |
| 71 | elif isinstance(result, Exception): |
| 72 | raise result |
| 73 | return result |
| 74 | |
| 75 | |
| 76 | class RequestTracker: |