jpayne@68: """
jpayne@68: Asynchronous progressbar decorator for iterators.
jpayne@68: Includes a default `range` iterator printing to `stderr`.
jpayne@68: 
jpayne@68: Usage:
jpayne@68: >>> from tqdm.asyncio import trange, tqdm
jpayne@68: >>> async for i in trange(10):
jpayne@68: ...     ...
jpayne@68: """
jpayne@68: import asyncio
jpayne@68: from sys import version_info
jpayne@68: 
jpayne@68: from .std import tqdm as std_tqdm
jpayne@68: 
jpayne@68: __author__ = {"github.com/": ["casperdcl"]}
jpayne@68: __all__ = ['tqdm_asyncio', 'tarange', 'tqdm', 'trange']
jpayne@68: 
jpayne@68: 
jpayne@68: class tqdm_asyncio(std_tqdm):
jpayne@68:     """
jpayne@68:     Asynchronous-friendly version of tqdm.
jpayne@68:     """
jpayne@68:     def __init__(self, iterable=None, *args, **kwargs):
jpayne@68:         super().__init__(iterable, *args, **kwargs)
jpayne@68:         self.iterable_awaitable = False
jpayne@68:         if iterable is not None:
jpayne@68:             if hasattr(iterable, "__anext__"):
jpayne@68:                 self.iterable_next = iterable.__anext__
jpayne@68:                 self.iterable_awaitable = True
jpayne@68:             elif hasattr(iterable, "__next__"):
jpayne@68:                 self.iterable_next = iterable.__next__
jpayne@68:             else:
jpayne@68:                 self.iterable_iterator = iter(iterable)
jpayne@68:                 self.iterable_next = self.iterable_iterator.__next__
jpayne@68: 
jpayne@68:     def __aiter__(self):
jpayne@68:         return self
jpayne@68: 
jpayne@68:     async def __anext__(self):
jpayne@68:         try:
jpayne@68:             if self.iterable_awaitable:
jpayne@68:                 res = await self.iterable_next()
jpayne@68:             else:
jpayne@68:                 res = self.iterable_next()
jpayne@68:             self.update()
jpayne@68:             return res
jpayne@68:         except StopIteration:
jpayne@68:             self.close()
jpayne@68:             raise StopAsyncIteration
jpayne@68:         except BaseException:
jpayne@68:             self.close()
jpayne@68:             raise
jpayne@68: 
jpayne@68:     def send(self, *args, **kwargs):
jpayne@68:         return self.iterable.send(*args, **kwargs)
jpayne@68: 
jpayne@68:     @classmethod
jpayne@68:     def as_completed(cls, fs, *, loop=None, timeout=None, total=None, **tqdm_kwargs):
jpayne@68:         """
jpayne@68:         Wrapper for `asyncio.as_completed`.
jpayne@68:         """
jpayne@68:         if total is None:
jpayne@68:             total = len(fs)
jpayne@68:         kwargs = {}
jpayne@68:         if version_info[:2] < (3, 10):
jpayne@68:             kwargs['loop'] = loop
jpayne@68:         yield from cls(asyncio.as_completed(fs, timeout=timeout, **kwargs),
jpayne@68:                        total=total, **tqdm_kwargs)
jpayne@68: 
jpayne@68:     @classmethod
jpayne@68:     async def gather(cls, *fs, loop=None, timeout=None, total=None, **tqdm_kwargs):
jpayne@68:         """
jpayne@68:         Wrapper for `asyncio.gather`.
jpayne@68:         """
jpayne@68:         async def wrap_awaitable(i, f):
jpayne@68:             return i, await f
jpayne@68: 
jpayne@68:         ifs = [wrap_awaitable(i, f) for i, f in enumerate(fs)]
jpayne@68:         res = [await f for f in cls.as_completed(ifs, loop=loop, timeout=timeout,
jpayne@68:                                                  total=total, **tqdm_kwargs)]
jpayne@68:         return [i for _, i in sorted(res)]
jpayne@68: 
jpayne@68: 
jpayne@68: def tarange(*args, **kwargs):
jpayne@68:     """
jpayne@68:     A shortcut for `tqdm.asyncio.tqdm(range(*args), **kwargs)`.
jpayne@68:     """
jpayne@68:     return tqdm_asyncio(range(*args), **kwargs)
jpayne@68: 
jpayne@68: 
jpayne@68: # Aliases
jpayne@68: tqdm = tqdm_asyncio
jpayne@68: trange = tarange