45 lines
1.3 KiB
Python
45 lines
1.3 KiB
Python
|
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)
|