Mercurial > repos > rliterman > csp2
annotate CSP2/CSP2_env/env-d9b9114564458d9d-741b3de822f2aaca6c6caa4325c4afce/lib/python3.8/site-packages/tqdm/dask.py @ 68:5028fdace37b
planemo upload commit 2e9511a184a1ca667c7be0c6321a36dc4e3d116d
| author | jpayne |
|---|---|
| date | Tue, 18 Mar 2025 16:23:26 -0400 |
| parents | |
| children |
| rev | line source |
|---|---|
| jpayne@68 | 1 from functools import partial |
| jpayne@68 | 2 |
| jpayne@68 | 3 from dask.callbacks import Callback |
| jpayne@68 | 4 |
| jpayne@68 | 5 from .auto import tqdm as tqdm_auto |
| jpayne@68 | 6 |
| jpayne@68 | 7 __author__ = {"github.com/": ["casperdcl"]} |
| jpayne@68 | 8 __all__ = ['TqdmCallback'] |
| jpayne@68 | 9 |
| jpayne@68 | 10 |
| jpayne@68 | 11 class TqdmCallback(Callback): |
| jpayne@68 | 12 """Dask callback for task progress.""" |
| jpayne@68 | 13 def __init__(self, start=None, pretask=None, tqdm_class=tqdm_auto, |
| jpayne@68 | 14 **tqdm_kwargs): |
| jpayne@68 | 15 """ |
| jpayne@68 | 16 Parameters |
| jpayne@68 | 17 ---------- |
| jpayne@68 | 18 tqdm_class : optional |
| jpayne@68 | 19 `tqdm` class to use for bars [default: `tqdm.auto.tqdm`]. |
| jpayne@68 | 20 tqdm_kwargs : optional |
| jpayne@68 | 21 Any other arguments used for all bars. |
| jpayne@68 | 22 """ |
| jpayne@68 | 23 super().__init__(start=start, pretask=pretask) |
| jpayne@68 | 24 if tqdm_kwargs: |
| jpayne@68 | 25 tqdm_class = partial(tqdm_class, **tqdm_kwargs) |
| jpayne@68 | 26 self.tqdm_class = tqdm_class |
| jpayne@68 | 27 |
| jpayne@68 | 28 def _start_state(self, _, state): |
| jpayne@68 | 29 self.pbar = self.tqdm_class(total=sum( |
| jpayne@68 | 30 len(state[k]) for k in ['ready', 'waiting', 'running', 'finished'])) |
| jpayne@68 | 31 |
| jpayne@68 | 32 def _posttask(self, *_, **__): |
| jpayne@68 | 33 self.pbar.update() |
| jpayne@68 | 34 |
| jpayne@68 | 35 def _finish(self, *_, **__): |
| jpayne@68 | 36 self.pbar.close() |
| jpayne@68 | 37 |
| jpayne@68 | 38 def display(self): |
| jpayne@68 | 39 """Displays in the current cell in Notebooks.""" |
| jpayne@68 | 40 container = getattr(self.bar, 'container', None) |
| jpayne@68 | 41 if container is None: |
| jpayne@68 | 42 return |
| jpayne@68 | 43 from .notebook import display |
| jpayne@68 | 44 display(container) |