codekingpro/portable-devtools

1from functools import partial2 3from dask.callbacks import Callback4 5from .auto import tqdm as tqdm_auto6 7__author__ = {"github.com/": ["casperdcl"]}8__all__ = ['TqdmCallback']9 10 11class TqdmCallback(Callback):12    """Dask callback for task progress."""13    def __init__(self, start=None, pretask=None, tqdm_class=tqdm_auto,14                 **tqdm_kwargs):15        """16        Parameters17        ----------18        tqdm_class  : optional19            `tqdm` class to use for bars [default: `tqdm.auto.tqdm`].20        tqdm_kwargs  : optional21            Any other arguments used for all bars.22        """23        super().__init__(start=start, pretask=pretask)24        if tqdm_kwargs:25            tqdm_class = partial(tqdm_class, **tqdm_kwargs)26        self.tqdm_class = tqdm_class27 28    def _start_state(self, _, state):29        self.pbar = self.tqdm_class(total=sum(30            len(state[k]) for k in ['ready', 'waiting', 'running', 'finished']))31 32    def _posttask(self, *_, **__):33        self.pbar.update()34 35    def _finish(self, *_, **__):36        self.pbar.close()37 38    def display(self):39        """Displays in the current cell in Notebooks."""40        container = getattr(self.bar, 'container', None)41        if container is None:42            return43        from .notebook import display44        display(container)45