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