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)