Merge pull request #1079 from tqdm/dask

casperdcl · web-flow · commit 2a82405da530 · 2021-03-05T13:48:10.000Z
diff --git a/.meta/.readme.rst b/.meta/.readme.rst
@@ -140,7 +140,7 @@ Changelog
 The list of all changes is available either on GitHub's Releases:
 |GitHub-Status|, on the
 `wiki <https://github.com/tqdm/tqdm/wiki/Releases>`__, or on the
-`website <https://tqdm.github.io/releases/>`__.
+`website <https://tqdm.github.io/releases>`__.
 
 
 Usage
@@ -460,7 +460,10 @@ Submodules
         """`rich.progress` version."""
 
     class tqdm.keras.TqdmCallback(keras.callbacks.Callback):
-        """`keras` callback for epoch and batch progress."""
+        """Keras callback for epoch and batch progress."""
+
+    class tqdm.dask.TqdmCallback(dask.callbacks.Callback):
+        """Dask callback for task progress."""
 
 
 ``contrib``
@@ -470,8 +473,8 @@ The ``tqdm.contrib`` package also contains experimental modules:
 
 - ``tqdm.contrib.itertools``: Thin wrappers around ``itertools``
 - ``tqdm.contrib.concurrent``: Thin wrappers around ``concurrent.futures``
-- ``tqdm.contrib.discord``: Posts to `Discord <https://discord.com/>`__ bots
-- ``tqdm.contrib.telegram``: Posts to `Telegram <https://telegram.org/>`__ bots
+- ``tqdm.contrib.discord``: Posts to `Discord <https://discord.com>`__ bots
+- ``tqdm.contrib.telegram``: Posts to `Telegram <https://telegram.org>`__ bots
 - ``tqdm.contrib.bells``: Automagically enables all optional features
 
   * ``auto``, ``pandas``, ``discord``, ``telegram``
@@ -829,6 +832,24 @@ A ``keras`` callback is also available:
 
     model.fit(..., verbose=0, callbacks=[TqdmCallback()])
 
+Dask Integration
+~~~~~~~~~~~~~~~~
+
+A ``dask`` callback is also available:
+
+.. code:: python
+
+    from tqdm.dask import TqdmCallback
+
+    with TqdmCallback(desc="compute"):
+        ...
+        arr.compute()
+
+    # or use callback globally
+    cb = TqdmCallback(desc="global")
+    cb.register()
+    arr.compute()
+
 IPython/Jupyter Integration
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
diff --git a/README.rst b/README.rst
@@ -140,7 +140,7 @@ Changelog
 The list of all changes is available either on GitHub's Releases:
 |GitHub-Status|, on the
 `wiki <https://github.com/tqdm/tqdm/wiki/Releases>`__, or on the
-`website <https://tqdm.github.io/releases/>`__.
+`website <https://tqdm.github.io/releases>`__.
 
 
 Usage
@@ -679,7 +679,10 @@ Submodules
         """`rich.progress` version."""
 
     class tqdm.keras.TqdmCallback(keras.callbacks.Callback):
-        """`keras` callback for epoch and batch progress."""
+        """Keras callback for epoch and batch progress."""
+
+    class tqdm.dask.TqdmCallback(dask.callbacks.Callback):
+        """Dask callback for task progress."""
 
 
 ``contrib``
@@ -689,8 +692,8 @@ The ``tqdm.contrib`` package also contains experimental modules:
 
 - ``tqdm.contrib.itertools``: Thin wrappers around ``itertools``
 - ``tqdm.contrib.concurrent``: Thin wrappers around ``concurrent.futures``
-- ``tqdm.contrib.discord``: Posts to `Discord <https://discord.com/>`__ bots
-- ``tqdm.contrib.telegram``: Posts to `Telegram <https://telegram.org/>`__ bots
+- ``tqdm.contrib.discord``: Posts to `Discord <https://discord.com>`__ bots
+- ``tqdm.contrib.telegram``: Posts to `Telegram <https://telegram.org>`__ bots
 - ``tqdm.contrib.bells``: Automagically enables all optional features
 
   * ``auto``, ``pandas``, ``discord``, ``telegram``
@@ -1048,6 +1051,24 @@ A ``keras`` callback is also available:
 
     model.fit(..., verbose=0, callbacks=[TqdmCallback()])
 
+Dask Integration
+~~~~~~~~~~~~~~~~
+
+A ``dask`` callback is also available:
+
+.. code:: python
+
+    from tqdm.dask import TqdmCallback
+
+    with TqdmCallback(desc="compute"):
+        ...
+        arr.compute()
+
+    # or use callback globally
+    cb = TqdmCallback(desc="global")
+    cb.register()
+    arr.compute()
+
 IPython/Jupyter Integration
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
diff --git a/benchmarks/README.md b/benchmarks/README.md
@@ -11,7 +11,7 @@ These benchmarks serve two purposes:
     - [`progressbar2`](https://pypi.org/project/progressbar2)
     - [`alive-progress`](https://pypi.org/project/alive-progress)
 
-Performance graphs are available at <https://tqdm.github.io/tqdm/>
+Performance graphs are available at <https://tqdm.github.io/tqdm>
 
 ## Running
 
diff --git a/environment.yml b/environment.yml
@@ -25,6 +25,7 @@ dependencies:
 - flake8-comprehensions
 - coverage
 # extras
+- dask               # dask
 - matplotlib         # gui
 - numpy              # pandas, keras, contrib.tenumerate
 - pandas
diff --git a/tests/tests_dask.py b/tests/tests_dask.py
@@ -0,0 +1,19 @@
+from __future__ import division
+
+from time import sleep
+
+from .tests_tqdm import importorskip, mark
+
+pytestmark = mark.slow
+
+
+def test_dask(capsys):
+    """Test tqdm.dask.TqdmCallback"""
+    ProgressBar = importorskip('tqdm.dask').TqdmCallback
+    dask = importorskip('dask')
+
+    schedule = [dask.delayed(sleep)(i / 10) for i in range(5)]
+    with ProgressBar():
+        dask.compute(schedule)
+    _, err = capsys.readouterr()
+    assert '5/5' in err
diff --git a/tox.ini b/tox.ini
@@ -27,6 +27,7 @@ passenv=TOXENV CI GITHUB_* CODECOV_* COVERALLS_* CODACY_* HOME
 deps=
     {[core]deps}
     cython
+    dask[delayed]
     matplotlib
     numpy
     pandas
diff --git a/tqdm/dask.py b/tqdm/dask.py
@@ -0,0 +1,46 @@
+from __future__ import absolute_import
+
+from functools import partial
+
+from dask.callbacks import Callback
+
+from .auto import tqdm as tqdm_auto
+
+__author__ = {"github.com/": ["casperdcl"]}
+__all__ = ['TqdmCallback']
+
+
+class TqdmCallback(Callback):
+    """Dask callback for task progress."""
+    def __init__(self, start=None, pretask=None, tqdm_class=tqdm_auto,
+                 **tqdm_kwargs):
+        """
+        Parameters
+        ----------
+        tqdm_class : optional
+            `tqdm` class to use for bars [default: `tqdm.auto.tqdm`].
+        tqdm_kwargs  : optional
+            Any other arguments used for all bars.
+        """
+        super(TqdmCallback, self).__init__(start=start, pretask=pretask)
+        if tqdm_kwargs:
+            tqdm_class = partial(tqdm_class, **tqdm_kwargs)
+        self.tqdm_class = tqdm_class
+
+    def _start_state(self, _, state):
+        self.pbar = self.tqdm_class(total=sum(
+            len(state[k]) for k in ['ready', 'waiting', 'running', 'finished']))
+
+    def _posttask(self, *_, **__):
+        self.pbar.update()
+
+    def _finish(self, *_, **__):
+        self.pbar.close()
+
+    def display(self):
+        """Displays in the current cell in Notebooks."""
+        container = getattr(self.bar, 'container', None)
+        if container is None:
+            return
+        from .notebook import display
+        display(container)
diff --git a/tqdm/gui.py b/tqdm/gui.py
@@ -24,12 +24,8 @@
 
 
 class tqdm_gui(std_tqdm):  # pragma: no cover
-    """
-    Experimental Matplotlib GUI version of tqdm!
-    """
-
+    """Experimental Matplotlib GUI version of tqdm!"""
     # TODO: @classmethod: write() on GUI?
-
     def __init__(self, *args, **kwargs):
         from collections import deque
 
diff --git a/tqdm/keras.py b/tqdm/keras.py
@@ -17,7 +17,7 @@
 
 
 class TqdmCallback(keras.callbacks.Callback):
-    """`keras` callback for epoch and batch progress"""
+    """Keras callback for epoch and batch progress."""
     @staticmethod
     def bar2callback(bar, pop=None, delta=(lambda logs: 1)):
         def callback(_, logs=None):
@@ -98,7 +98,7 @@ def on_train_end(self, *_, **__):
         self.epoch_bar.close()
 
     def display(self):
-        """displays in the current cell in Notebooks"""
+        """Displays in the current cell in Notebooks."""
         container = getattr(self.epoch_bar, 'container', None)
         if container is None:
             return
diff --git a/tqdm/rich.py b/tqdm/rich.py
@@ -74,12 +74,8 @@ def render(self, task):
 
 
 class tqdm_rich(std_tqdm):  # pragma: no cover
-    """
-    Experimental rich.progress GUI version of tqdm!
-    """
-
+    """Experimental rich.progress GUI version of tqdm!"""
     # TODO: @classmethod: write()?
-
     def __init__(self, *args, **kwargs):
         """
         This class accepts the following parameters *in addition* to
