Current File : //opt/imunify360/venv/lib/python3.11/site-packages/defence360agent/utils/async_utils.py
from typing import List, Union, Tuple
import asyncio


class AsyncIterate:  # not AsyncIterable because python use this name already
    def __init__(self, data: Union[List, Tuple]):
        self.queue = iter(data)

    def __aiter__(self):
        return self

    async def __anext__(self):
        data = await self.fetch_data()
        if data is not None:
            return data
        else:
            raise StopAsyncIteration

    async def fetch_data(self):
        try:
            item = next(self.queue)
        except StopIteration:
            item = None
        return item


async def gather(*tasks: List) -> AsyncIterate:
    results = await asyncio.gather(*tasks)
    return AsyncIterate(results)