Voice et bot modif

venv/lib/python3.12/site-packages/tqdm/contrib/itertools.py

@@ -6,30 +6,96 @@ import itertools
 from ..auto import tqdm as tqdm_auto
 
 __author__ = {"github.com/": ["casperdcl"]}
-__all__ = ['product']
+__all__ = [
+    'chain', 'product', 'permutations', 'combinations', 'combinations_with_replacement', 'batched']
 
 
-def product(*iterables, **tqdm_kwargs):
-    """
-    Equivalent of `itertools.product`.
+def chain(*iterables, total=None, tqdm_class=tqdm_auto, **kwargs):
+    """Equivalent of `itertools.chain`."""
+    if total is None:
+        try:
+            total = sum(map(len, iterables))
+        except (TypeError, AttributeError):
+            pass
+    return tqdm_class(itertools.chain(*iterables), total=total, **kwargs)
 
-    Parameters
-    ----------
-    tqdm_class  : [default: tqdm.auto.tqdm].
-    """
-    kwargs = tqdm_kwargs.copy()
-    tqdm_class = kwargs.pop("tqdm_class", tqdm_auto)
-    try:
-        lens = list(map(len, iterables))
-    except TypeError:
-        total = None
-    else:
-        total = 1
-        for i in lens:
-            total *= i
-        kwargs.setdefault("total", total)
-    with tqdm_class(**kwargs) as t:
-        it = itertools.product(*iterables)
-        for i in it:
-            yield i
-            t.update()
+
+def product(*iterables, repeat=1, total=None, tqdm_class=tqdm_auto, **kwargs):
+    """Equivalent of `itertools.product`."""
+    if total is None:
+        try:
+            lens = list(map(len, iterables))
+        except (TypeError, AttributeError):
+            pass
+        else:  # py>=3.8: math.prod(lens) ** repeat
+            total = 1
+            for i in lens:
+                total *= i
+            total **= repeat
+    for i in tqdm_class(itertools.product(*iterables, repeat=repeat), total=total, **kwargs):
+        yield i
+
+
+def permutations(iterable, r=None, total=None, tqdm_class=tqdm_auto, **kwargs):
+    """Equivalent of `itertools.permutations`."""
+    if total is None:
+        try:
+            n = len(iterable)
+        except (TypeError, AttributeError):
+            pass
+        else:
+            r = n if r is None else r
+            if r > n:
+                total = 0
+            else:  # py>=3.8: math.perm(n, r)
+                total = 1
+                for i in range(n, n-r, -1):
+                    total *= i
+    return tqdm_class(itertools.permutations(iterable, r), total=total, **kwargs)
+
+
+def combinations(iterable, r, total=None, tqdm_class=tqdm_auto, **kwargs):
+    """Equivalent of `itertools.combinations`."""
+    if total is None:
+        try:
+            n = len(iterable)
+        except (TypeError, AttributeError):
+            pass
+        else:
+            if r > n:
+                total = 0
+            else:  # py>=3.8: math.comb(n, r)
+                total = 1
+                for i in range(n, n-r, -1):
+                    total *= i
+                for i in range(1, r+1):
+                    total //= i
+    return tqdm_class(itertools.combinations(iterable, r), total=total, **kwargs)
+
+
+def combinations_with_replacement(iterable, r, total=None, tqdm_class=tqdm_auto, **kwargs):
+    """Equivalent of `itertools.combinations_with_replacement`."""
+    if total is None:
+        try:
+            n = len(iterable)
+        except (TypeError, AttributeError):
+            pass
+        else:
+            total = 1
+            for i in range(n+r-1, n-1, -1):
+                total *= i
+            for i in range(1, r+1):
+                total //= i
+    return tqdm_class(itertools.combinations_with_replacement(iterable, r), total=total, **kwargs)
+
+
+def batched(iterable, n, total=None, tqdm_class=tqdm_auto, **kwargs):
+    """Equivalent of `itertools.batched`."""
+    if total is None:
+        try:
+            total = len(iterable)
+        except (TypeError, AttributeError):
+            pass
+    return tqdm_class(itertools.batched(iterable, n), unit_scale=n,
+                      total=(total+n-1) // n if total is not None else None,
+                      **kwargs)