TQDMDaskProgressBar: tqdm for dask 1.2.2

tqdm_dask_progressbar.py

from dask.callbacks import Callback
from tqdm.auto import tqdm

class TQDMDaskProgressBar(Callback, object):
    """
    A tqdm progress bar for dask.
    Usage:
        ```
        with TQDMDaskProgressBar():
            da.compute()
        ```
    Author : wassname, jmcarpenter2
    Source : https://gist.github.com/wassname/1837d0365247430e02abda41f0e7f184
    See: http://dask.pydata.org/en/latest/diagnostics-local.html?highlight=progress
    """

    def __init__(self, start=None, start_state=None, pretask=None, posttask=None, finish=None, **kwargs):
        super(TQDMDaskProgressBar, self).__init__(
            start=start, start_state=start_state, pretask=pretask, posttask=posttask, finish=finish
        )
        self.tqdm_args = kwargs
        self.states = ["ready", "waiting", "running", "finished"]

    def _start_state(self, dsk, state):
        self._tqdm = tqdm(total=sum(len(state[k]) for k in self.states), **self.tqdm_args)

    def _pretask(self, key, dsk, state):
        self.state = state
        self._tqdm.desc = '{:20.20s}'.format(str(key))

    def _posttask(self, key, result, dsk, state, worker_id):
        self._tqdm.update(1)

    def _finish(self, dsk, state, errored):
        self._tqdm.close()

also see improvements https://github.com/jmcarpenter2/swifter/blob/master/swifter/tqdm_dask_progressbar.py