Add progress bar with tqdm.

Author: mancellin

labelled_functions/decorators.py

@@ -50,6 +50,25 @@ def timed_func(*args, **kwargs):
     return timed_func
 
 
+def with_progress_bar(lab_f, total=None):
+    """Add a tqdm object to count calls and display a progress bar."""
+    from tqdm import tqdm
+    bar = tqdm(total=total, unit="calls")
+
+    def add_call_counter(f):
+        def exec_with_call_counter(*args, **kwargs):
+            out = f(*args, **kwargs)
+            bar.update(1)
+            return out
+        return exec_with_call_counter
+
+    new_lab_f = copy(lab_f)
+    new_lab_f.bar = bar
+    new_lab_f.function = add_call_counter(lab_f.function)
+
+    return new_lab_f
+
+
 def decorate(func, decorator):
     new_func = copy(func)
     new_func.function = decorator(func.function)

labelled_functions/maps.py

@@ -8,14 +8,16 @@
 import xarray as xr
 
 from .labels import label
-from .decorators import keeping_inputs
+from .decorators import keeping_inputs, with_progress_bar
 
 
 # API
 
-def pandas_map(f, *args, n_jobs=1, **kwargs):
+def pandas_map(f, *args, progress_bar=False, n_jobs=1, **kwargs):
     f = label(f)
     dict_of_lists = _preprocess_map_inputs(f.input_names, args, kwargs)
+    if progress_bar:
+        f = with_progress_bar(f, total=len(any_value(dict_of_lists)))
     if n_jobs == 1:
         data = list(lmap(keeping_inputs(f), **dict_of_lists))
     else:
@@ -25,9 +27,11 @@ def pandas_map(f, *args, n_jobs=1, **kwargs):
     return _set_index(f.input_names, data)
 
 
-def pandas_cartesian_product(f, *args, n_jobs=1, **kwargs):
+def pandas_cartesian_product(f, *args, progress_bar=False, n_jobs=1, **kwargs):
     f = label(f)
     dict_of_lists = _preprocess_map_inputs(f.input_names, args, kwargs)
+    if progress_bar:
+        f = with_progress_bar(f, total=len(any_value(dict_of_lists)))
     if n_jobs == 1:
         data = list(lcartesianmap(keeping_inputs(f), **dict_of_lists))
     else:
@@ -91,3 +95,7 @@ def _set_index(indices, data):
         return data.set_index(indices)
     else:
         return data
+
+def any_value(d):
+    """Returns one of the values of a dict."""
+    return next(iter(d.values()))

setup.py

@@ -22,5 +22,6 @@
         'toolz',
         'parso',
         'joblib',
+        'tqdm',
     ],
 )

test/test_decorators.py

@@ -3,8 +3,11 @@
 
 import pytest
 
+from time import sleep
+
 from labelled_functions import label, pipeline
-from labelled_functions.decorators import keeping_inputs, decorate, cache
+from labelled_functions.maps import pandas_map
+from labelled_functions.decorators import keeping_inputs, with_progress_bar, decorate, cache
 from example_functions import *
 
 
@@ -40,3 +43,16 @@ def test_cache():
     a = pipe(length=1.0)
     b = pipe(length=1.0)
     assert a == b
+
+
+def test_progress_bar(capfd):
+    def wait(dt):
+        sleep(dt)
+        output = 1
+        return output
+
+    pandas_map(wait, dt=[0.01]*10, progress_bar=True)
+
+    out, err = capfd.readouterr()
+    assert "10/10" in err
+