Upstream tqdm/tqdm state of the touched files before bdd840a9a1

tests/tests_asyncio.py

@@ -0,0 +1,184 @@
+"""Tests `tqdm.asyncio`."""
+import asyncio
+from functools import partial
+from sys import platform
+from time import time
+
+from tqdm.asyncio import tarange, tqdm_asyncio
+
+from .tests_tqdm import StringIO, closing, mark
+
+tqdm = partial(tqdm_asyncio, miniters=0, mininterval=0)
+trange = partial(tarange, miniters=0, mininterval=0)
+as_completed = partial(tqdm_asyncio.as_completed, miniters=0, mininterval=0)
+gather = partial(tqdm_asyncio.gather, miniters=0, mininterval=0)
+
+
+def count(start=0, step=1):
+    i = start
+    while True:
+        new_start = yield i
+        if new_start is None:
+            i += step
+        else:
+            i = new_start
+
+
+async def acount(*args, **kwargs):
+    for i in count(*args, **kwargs):
+        yield i
+
+
+@mark.asyncio
+async def test_break():
+    """Test asyncio break"""
+    pbar = tqdm(count())
+    async for _ in pbar:
+        break
+    pbar.close()
+
+
+@mark.asyncio
+async def test_generators(capsys):
+    """Test asyncio generators"""
+    with tqdm(count(), desc="counter") as pbar:
+        async for i in pbar:
+            if i >= 8:
+                break
+    _, err = capsys.readouterr()
+    assert '9it' in err
+
+    acounter = acount()
+    try:
+        with tqdm(acounter, desc="async_counter") as pbar:
+            async for i in pbar:
+                if i >= 8:
+                    break
+    finally:
+        await acounter.aclose()
+    _, err = capsys.readouterr()
+    assert '9it' in err
+
+
+@mark.asyncio
+async def test_range():
+    """Test asyncio range"""
+    with closing(StringIO()) as our_file:
+        async for _ in tqdm(range(9), desc="range", file=our_file):
+            pass
+        assert '9/9' in our_file.getvalue()
+        our_file.seek(0)
+        our_file.truncate()
+
+        async for _ in trange(9, desc="trange", file=our_file):
+            pass
+        assert '9/9' in our_file.getvalue()
+
+
+@mark.asyncio
+async def test_nested():
+    """Test asyncio nested"""
+    with closing(StringIO()) as our_file:
+        async for _ in tqdm(trange(9, desc="inner", file=our_file),
+                            desc="outer", file=our_file):
+            pass
+        assert 'inner: 100%' in our_file.getvalue()
+        assert 'outer: 100%' in our_file.getvalue()
+
+
+@mark.asyncio
+async def test_coroutines():
+    """Test asyncio coroutine.send"""
+    with closing(StringIO()) as our_file:
+        with tqdm(count(), file=our_file) as pbar:
+            async for i in pbar:
+                if i == 9:
+                    pbar.send(-10)
+                elif i < 0:
+                    assert i == -9
+                    break
+        assert '10it' in our_file.getvalue()
+
+
+@mark.slow
+@mark.asyncio
+@mark.parametrize("tol", [0.2 if platform.startswith("darwin") else 0.1])
+async def test_as_completed(capsys, tol):
+    """Test asyncio as_completed"""
+    for retry in range(3):
+        t = time()
+        skew = time() - t
+        for i in as_completed([asyncio.sleep(0.01 * i) for i in range(30, 0, -1)]):
+            await i
+        t = time() - t - 2 * skew
+        try:
+            assert 0.3 * (1 - tol) < t < 0.3 * (1 + tol), t
+            _, err = capsys.readouterr()
+            assert '30/30' in err
+        except AssertionError:
+            if retry == 2:
+                raise
+
+
+async def double(i):
+    return i * 2
+
+
+@mark.asyncio
+async def test_gather(capsys):
+    """Test asyncio gather"""
+    res = await gather(*map(double, range(30)))
+    _, err = capsys.readouterr()
+    assert '30/30' in err
+    assert res == list(range(0, 30 * 2, 2))
+
+
+async def raise_exc(i):
+    if i == 1:
+        raise ValueError("test")
+    return i * 2
+
+
+@mark.asyncio
+async def test_gather_exceptions(capsys):
+    """Test asyncio gather return_exceptions"""
+    res = await gather(*map(raise_exc, range(3)), return_exceptions=True)
+    _, err = capsys.readouterr()
+    assert '3/3' in err
+    assert isinstance(res[1], ValueError)
+    assert res[0] == 0
+    assert res[2] == 4
+
+
+class AIterable:
+    """Has `__aiter__` but not `__anext__`"""
+    def __init__(self, items):
+        self._items = items
+
+    def __aiter__(self):
+        return AIterator(self._items)
+
+
+class AIterator:
+    def __init__(self, items):
+        self._items = iter(items)
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            return next(self._items)
+        except StopIteration:
+            raise StopAsyncIteration
+
+
+@mark.asyncio
+async def test_async_iterable(capsys):
+    """Test asyncio with async iterable (only __aiter__, no __anext__)"""
+    result = []
+    async for i in tqdm(AIterable(range(9)), desc="aiterable"):
+        result.append(i)
+    assert result == list(range(9))
+    _, err = capsys.readouterr()
+    assert '9it' in err

tqdm/asyncio.py

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