Use total kwarg as tqdm length if provided

swansonk14 · Jun 14, 2021 · 6a755c8 · 6a755c8
1 parent a141b89
commit 6a755c8
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/p_tqdm/p_tqdm.py b/p_tqdm/p_tqdm.py
@@ -38,8 +38,9 @@ def _parallel(ordered: bool, function: Callable, *iterables: Iterable, **kwargs:
     elif type(num_cpus) == float:
         num_cpus = int(round(num_cpus * cpu_count()))
 
-    # Determine length of tqdm (equal to length of shortest iterable)
-    length = min(len(iterable) for iterable in iterables if isinstance(iterable, Sized))
+    # Determine length of tqdm (equal to length of shortest iterable or total kwarg)
+    total = kwargs.pop('total', None)
+    length = total or min(len(iterable) for iterable in iterables if isinstance(iterable, Sized))
 
     # Create parallel generator
     map_type = 'imap' if ordered else 'uimap'