12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061 |
- """
- Tests for `tqdm.contrib`.
- """
- import pytest
- from tqdm import tqdm
- from tqdm.contrib import tenumerate, tmap, tzip
- from .tests_tqdm import StringIO, closing, importorskip
- def incr(x):
- """Dummy function"""
- return x + 1
- @pytest.mark.parametrize("tqdm_kwargs", [{}, {"tqdm_class": tqdm}])
- def test_enumerate(tqdm_kwargs):
- """Test contrib.tenumerate"""
- with closing(StringIO()) as our_file:
- a = range(9)
- assert list(tenumerate(a, file=our_file, **tqdm_kwargs)) == list(enumerate(a))
- assert list(tenumerate(a, 42, file=our_file, **tqdm_kwargs)) == list(
- enumerate(a, 42)
- )
- with closing(StringIO()) as our_file:
- _ = list(tenumerate(iter(a), file=our_file, **tqdm_kwargs))
- assert "100%" not in our_file.getvalue()
- with closing(StringIO()) as our_file:
- _ = list(tenumerate(iter(a), file=our_file, total=len(a), **tqdm_kwargs))
- assert "100%" in our_file.getvalue()
- def test_enumerate_numpy():
- """Test contrib.tenumerate(numpy.ndarray)"""
- np = importorskip("numpy")
- with closing(StringIO()) as our_file:
- a = np.random.random((42, 7))
- assert list(tenumerate(a, file=our_file)) == list(np.ndenumerate(a))
- @pytest.mark.parametrize("tqdm_kwargs", [{}, {"tqdm_class": tqdm}])
- def test_zip(tqdm_kwargs):
- """Test contrib.tzip"""
- with closing(StringIO()) as our_file:
- a = range(9)
- b = [i + 1 for i in a]
- gen = tzip(a, b, file=our_file, **tqdm_kwargs)
- assert gen != list(zip(a, b))
- assert list(gen) == list(zip(a, b))
- @pytest.mark.parametrize("tqdm_kwargs", [{}, {"tqdm_class": tqdm}])
- def test_map(tqdm_kwargs):
- """Test contrib.tmap"""
- with closing(StringIO()) as our_file:
- a = range(9)
- b = [i + 1 for i in a]
- gen = tmap(lambda x: x + 1, a, file=our_file, **tqdm_kwargs)
- assert gen != b
- assert list(gen) == b
|