Merge pull request #35 from fcoclavero/feature/total-kwarg

Use the total kwarg as the tqdm length if provided.
swansonk14 · Apr 19, 2022 · f7414a9 · f7414a9
2 parents edd6325 + 6a755c8
commit f7414a9
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/.gitignore b/.gitignore
@@ -6,3 +6,5 @@ p_tqdm.egg-info
 __pycache__
 *.pyc
 .idea
+venv
+.venv
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'
-Original file line number
+Diff line change
@@ Expand Up / @@ -6,3 +6,5 @@ p_tqdm.egg-info @@
     __pycache__
     *.pyc
     .idea
+    venv
+    .venv