| 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) | |