1""" 2Asynchronous progressbar decorator for iterators. 3Includes a default `range` iterator printing to `stderr`. 4 5Usage: 6>>> from tqdm.asyncio import trange, tqdm 7>>> async for i in trange(10): 8... ... 9""" 10import asyncio 11from sys import version_info 12 13from .std import tqdm as std_tqdm 14 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 (Python 3.6+). 22 """ 23 def __init__(self, iterable=None, *args, **kwargs): 24 super(tqdm_asyncio, self).__init__(iterable, *args, **kwargs) 25 self.iterable_awaitable = False 26 if iterable is not None: 27 if hasattr(iterable, "__anext__"): 28 self.iterable_next = iterable.__anext__ 29 self.iterable_awaitable = True 30 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 self 38 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 res 47 except StopIteration: 48 self.close() 49 raise StopAsyncIteration 50 except BaseException: 51 self.close() 52 raise 53 54 def send(self, *args, **kwargs): 55 return self.iterable.send(*args, **kwargs) 56 57 @classmethod 58 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'] = loop 67 yield from cls(asyncio.as_completed(fs, timeout=timeout, **kwargs), 68 total=total, **tqdm_kwargs) 69 70 @classmethod 71 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 f 77 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# Aliases 92tqdm = tqdm_asyncio 93trange = tarange 94