From fcece3c437d3dd5576d38650626425f9d87e1b17 Mon Sep 17 00:00:00 2001 From: Igor Dikarev Date: Mon, 12 Feb 2024 01:41:22 +0300 Subject: [PATCH] add keyword 'repeat' for tqdm.contrib.itertools.product --- tests/tests_itertools.py | 2 ++ tqdm/contrib/itertools.py | 4 ++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/tests_itertools.py b/tests/tests_itertools.py index bfb6eb2cd..0c030e751 100644 --- a/tests/tests_itertools.py +++ b/tests/tests_itertools.py @@ -24,3 +24,5 @@ def test_product(): assert list(product(a, a[::-1], file=our_file)) == list(it.product(a, a[::-1])) assert list(product(a, NoLenIter(a), file=our_file)) == list(it.product(a, NoLenIter(a))) + + assert list(product(a, repeat=2, file=our_file)) == list(it.product(a, repeat=2)) diff --git a/tqdm/contrib/itertools.py b/tqdm/contrib/itertools.py index e67651a41..2d3a7f194 100644 --- a/tqdm/contrib/itertools.py +++ b/tqdm/contrib/itertools.py @@ -9,7 +9,7 @@ __all__ = ['product'] -def product(*iterables, **tqdm_kwargs): +def product(*iterables, repeat=1, **tqdm_kwargs): """ Equivalent of `itertools.product`. @@ -29,7 +29,7 @@ def product(*iterables, **tqdm_kwargs): total *= i kwargs.setdefault("total", total) with tqdm_class(**kwargs) as t: - it = itertools.product(*iterables) + it = itertools.product(*iterables, repeat=repeat) for i in it: yield i t.update()