1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
|
from __future__ import absolute_import
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)
|
