add a progress_bar for loading a BalancingLearner

Author: basnijholt

adaptive/learner/balancing_learner.py

@@ -377,7 +377,7 @@ def save(self, fname, compress=True):
             for l in self.learners:
                 l.save(fname(l), compress=compress)
 
-    def load(self, fname, compress=True):
+    def load(self, fname, compress=True, with_progress_bar=False):
         """Load the data of the child learners from pickle files
         in a directory.
 
@@ -389,16 +389,30 @@ def load(self, fname, compress=True):
         compress : bool, default True
             If the data is compressed when saved, one must load it
             with compression too.
+        with_progress_bar : bool, default False
+            Display a progress bar using `tqdm`.
 
         Example
         -------
         See the example in the `BalancingLearner.save` doc-string.
         """
+        def progress(seq):
+            if not with_progress_bar:
+                return seq
+            else:
+                from adaptive.notebook_integration import in_ipynb
+                if in_ipynb():
+                    from tqdm import tqdm_notebook
+                    return tqdm_notebook(list(seq))
+                else:
+                    from tqdm import tqdm
+                    return tqdm(list(seq))
+
         if isinstance(fname, Iterable):
-            for l, _fname in zip(self.learners, fname):
+            for l, _fname in progress(zip(self.learners, fname)):
                 l.load(_fname, compress=compress)
         else:
-            for l in self.learners:
+            for l in progress(self.learners):
                 l.load(fname(l), compress=compress)
 
     def _get_data(self):

environment.yml

@@ -16,3 +16,4 @@ dependencies:
   - ipywidgets
   - scikit-optimize
   - plotly
+  - tqdm

setup.py

@@ -36,6 +36,7 @@ def get_version_and_cmdclass(package_name):
         "bokeh",
         "matplotlib",
         "plotly",
+        "tqdm",
     ]
 }