codekingpro/portable-devtools

1"""2Asynchronous progressbar decorator for iterators.3Includes a default `range` iterator printing to `stderr`.4 5Usage:6>>> from tqdm.asyncio import trange, tqdm7>>> async for i in trange(10):8...     ...9"""10import asyncio11from sys import version_info12 13from .std import tqdm as std_tqdm14 15__author__ = {"github.com/": ["casperdcl"]}16__all__ = ['tqdm_asyncio', 'tarange', 'tqdm', 'trange']17 18 19class tqdm_asyncio(std_tqdm):20    """21    Asynchronous-friendly version of tqdm.22    """23    def __init__(self, iterable=None, *args, **kwargs):24        super().__init__(iterable, *args, **kwargs)25        self.iterable_awaitable = False26        if iterable is not None:27            if hasattr(iterable, "__anext__"):28                self.iterable_next = iterable.__anext__29                self.iterable_awaitable = True30            elif hasattr(iterable, "__next__"):31                self.iterable_next = iterable.__next__32            else:33                self.iterable_iterator = iter(iterable)34                self.iterable_next = self.iterable_iterator.__next__35 36    def __aiter__(self):37        return self38 39    async def __anext__(self):40        try:41            if self.iterable_awaitable:42                res = await self.iterable_next()43            else:44                res = self.iterable_next()45            self.update()46            return res47        except StopIteration:48            self.close()49            raise StopAsyncIteration50        except BaseException:51            self.close()52            raise53 54    def send(self, *args, **kwargs):55        return self.iterable.send(*args, **kwargs)56 57    @classmethod58    def as_completed(cls, fs, *, loop=None, timeout=None, total=None, **tqdm_kwargs):59        """60        Wrapper for `asyncio.as_completed`.61        """62        if total is None:63            total = len(fs)64        kwargs = {}65        if version_info[:2] < (3, 10):66            kwargs['loop'] = loop67        yield from cls(asyncio.as_completed(fs, timeout=timeout, **kwargs),68                       total=total, **tqdm_kwargs)69 70    @classmethod71    async def gather(cls, *fs, loop=None, timeout=None, total=None, **tqdm_kwargs):72        """73        Wrapper for `asyncio.gather`.74        """75        async def wrap_awaitable(i, f):76            return i, await f77 78        ifs = [wrap_awaitable(i, f) for i, f in enumerate(fs)]79        res = [await f for f in cls.as_completed(ifs, loop=loop, timeout=timeout,80                                                 total=total, **tqdm_kwargs)]81        return [i for _, i in sorted(res)]82 83 84def tarange(*args, **kwargs):85    """86    A shortcut for `tqdm.asyncio.tqdm(range(*args), **kwargs)`.87    """88    return tqdm_asyncio(range(*args), **kwargs)89 90 91# Aliases92tqdm = tqdm_asyncio93trange = tarange94