asyncio.py 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. """
  2. Asynchronous progressbar decorator for iterators.
  3. Includes a default `range` iterator printing to `stderr`.
  4. Usage:
  5. >>> from tqdm.asyncio import trange, tqdm
  6. >>> async for i in trange(10):
  7. ... ...
  8. """
  9. from .std import tqdm as std_tqdm
  10. import asyncio
  11. __author__ = {"github.com/": ["casperdcl"]}
  12. __all__ = ['tqdm_asyncio', 'tarange', 'tqdm', 'trange']
  13. class tqdm_asyncio(std_tqdm):
  14. """
  15. Asynchronous-friendly version of tqdm (Python 3.5+).
  16. """
  17. def __init__(self, iterable=None, *args, **kwargs):
  18. super(tqdm_asyncio, self).__init__(iterable, *args, **kwargs)
  19. self.iterable_awaitable = False
  20. if iterable is not None:
  21. if hasattr(iterable, "__anext__"):
  22. self.iterable_next = iterable.__anext__
  23. self.iterable_awaitable = True
  24. elif hasattr(iterable, "__next__"):
  25. self.iterable_next = iterable.__next__
  26. else:
  27. self.iterable_iterator = iter(iterable)
  28. self.iterable_next = self.iterable_iterator.__next__
  29. def __aiter__(self):
  30. return self
  31. async def __anext__(self):
  32. try:
  33. if self.iterable_awaitable:
  34. res = await self.iterable_next()
  35. else:
  36. res = self.iterable_next()
  37. self.update()
  38. return res
  39. except StopIteration:
  40. self.close()
  41. raise StopAsyncIteration
  42. except:
  43. self.close()
  44. raise
  45. def send(self, *args, **kwargs):
  46. return self.iterable.send(*args, **kwargs)
  47. @classmethod
  48. def as_completed(cls, fs, *, loop=None, timeout=None, total=None,
  49. **tqdm_kwargs):
  50. """
  51. Wrapper for `asyncio.as_completed`.
  52. """
  53. if total is None:
  54. total = len(fs)
  55. yield from cls(asyncio.as_completed(fs, loop=loop, timeout=timeout),
  56. total=total, **tqdm_kwargs)
  57. def tarange(*args, **kwargs):
  58. """
  59. A shortcut for `tqdm.asyncio.tqdm(range(*args), **kwargs)`.
  60. """
  61. return tqdm_asyncio(range(*args), **kwargs)
  62. # Aliases
  63. tqdm = tqdm_asyncio
  64. trange = tarange