diff options
Diffstat (limited to '.venv/lib/python3.12/site-packages/tqdm/dask.py')
-rw-r--r-- | .venv/lib/python3.12/site-packages/tqdm/dask.py | 44 |
1 files changed, 44 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/tqdm/dask.py b/.venv/lib/python3.12/site-packages/tqdm/dask.py new file mode 100644 index 00000000..57f1b668 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/tqdm/dask.py @@ -0,0 +1,44 @@ +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().__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) |