Add dask based progress bar

esm_analysis/Reader.py

@@ -32,6 +32,8 @@
 def progress_bar(*futures, **kwargs):
     """Connect dask futures to tqdm progressbar."""
 
+    from .utils import (ProgressBar, NotebookProgress)
+
     notebook = kwargs.pop('notebook', None)
     multi = kwargs.pop('multi', True)
     complete = kwargs.pop('complete', True)
@@ -45,8 +47,10 @@ def progress_bar(*futures, **kwargs):
     if notebook is None:
         notebook = is_kernel()  # often but not always correct assumption
 
-    progress = tqdm.tqdm_notebook if notebook else tqdm.tqdm
-    _ = list(progress(as_completed(futures), **kwargs))
+    #progress = tqdm.tqdm_notebook if notebook else tqdm.tqdm
+    #_ = list(progress(as_completed(futures), **kwargs))
+    progress = ProgressBar if notebook else NotebookProgress
+    return progress(futures, **kwargs)
 
 
 class _BaseVariables(dict):

esm_analysis/tests/test_reader.py

@@ -107,7 +107,10 @@ def test_client(mock_run, mock_client):
 
 def test_progress_bar(mock_client):
     future = mock_client.submit(lambda: 2*6)
-    progress_bar(future)
+    progress_bar(future, notebook=False)
+    assert future.result() == 12
+    future = mock_client.submit(lambda: 2*6)
+    progress_bar(future, notebook=True)
     assert future.result() == 12
 
 def test_apply_function(mock_run, mock_client):

esm_analysis/utils.py

@@ -0,0 +1,59 @@
+import tqdm
+from tornado.ioloop import IOLoop
+from distributed.utils import LoopRunner, is_kernel
+from distributed.client import futures_of
+from distributed.diagnostics.progressbar import ProgressBar
+
+class ProgressBar(ProgressBar):
+    def __init__(
+        self,
+        keys,
+        scheduler=None,
+        interval="100ms",
+        loop=None,
+        complete=True,
+        start=True,
+        **tqdm_kwargs
+    ):
+        super(ProgressBar, self).__init__(
+            keys, scheduler, interval, complete)
+        self.tqdm = tqdm.tqdm(keys, **tqdm_kwargs)
+        self.loop = loop or IOLoop()
+
+        if start:
+            loop_runner = LoopRunner(self.loop)
+            loop_runner.run_sync(self.listen)
+
+    def _draw_bar(self, remaining, all, **kwargs):
+        update_ct = (all - remaining) - self.tqdm.n
+        self.tqdm.update(update_ct)
+
+    def _draw_stop(self, **kwargs):
+        self.tqdm.close()
+
+class NotebookProgress(ProgressBar):
+    def __init__(
+        self,
+        keys,
+        scheduler=None,
+        interval="100ms",
+        loop=None,
+        complete=True,
+        start=True,
+        **tqdm_kwargs
+    ):
+        super(NotebookProgress, self).__init__(
+            keys, scheduler, interval, complete)
+        self.tqdm = tqdm.tqdm_notebook(keys, **tqdm_kwargs)
+        self.loop = loop or IOLoop()
+
+        if start:
+            loop_runner = LoopRunner(self.loop)
+            loop_runner.run_sync(self.listen)
+
+    def _draw_bar(self, remaining, all, **kwargs):
+        update_ct = (all - remaining) - self.tqdm.n
+        self.tqdm.update(update_ct)
+
+    def _draw_stop(self, **kwargs):
+        self.tqdm.close()