summaryrefslogtreecommitdiffhomepage
path: root/libs/tqdm/tests/tests_itertools.py
blob: c55e07db87659b886b976dffc6f8086d7659bcee (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
"""
Tests for `tqdm.contrib.itertools`.
"""
from tqdm.contrib.itertools import product
from tests_tqdm import with_setup, pretest, posttest, StringIO, closing
import itertools


class NoLenIter(object):
    def __init__(self, iterable):
        self._it = iterable

    def __iter__(self):
        for i in self._it:
            yield i


@with_setup(pretest, posttest)
def test_product():
    """Test contrib.itertools.product"""
    with closing(StringIO()) as our_file:
        a = range(9)
        assert list(product(a, a[::-1], file=our_file)) == \
            list(itertools.product(a, a[::-1]))

        assert list(product(a, NoLenIter(a), file=our_file)) == \
            list(itertools.product(a, NoLenIter(a)))