commit e80bb3a1e7095a2a63ad3d1f20b51ea774709690 Author: upstream Date: Thu Sep 3 14:12:03 2026 +0200 Upstream tqdm/tqdm state of the touched files before bdd840a9a1 diff --git a/tests/tests_asyncio.py b/tests/tests_asyncio.py new file mode 100644 index 0000000..5311701 --- /dev/null +++ b/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 diff --git a/tqdm/asyncio.py b/tqdm/asyncio.py new file mode 100644 index 0000000..ac5c2e6 --- /dev/null +++ b/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