Mercurial > repos > rliterman > csp2

annotate CSP2/CSP2_env/env-d9b9114564458d9d-741b3de822f2aaca6c6caa4325c4afce/lib/python3.8/site-packages/tqdm/rich.py @ 68:5028fdace37b

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
planemo upload commit 2e9511a184a1ca667c7be0c6321a36dc4e3d116d
author jpayne
date Tue, 18 Mar 2025 16:23:26 -0400
parents
children
rev   line source
jpayne@68 1 """
jpayne@68 2 `rich.progress` decorator for iterators.
jpayne@68 3
jpayne@68 4 Usage:
jpayne@68 5 >>> from tqdm.rich import trange, tqdm
jpayne@68 6 >>> for i in trange(10):
jpayne@68 7 ... ...
jpayne@68 8 """
jpayne@68 9 from warnings import warn
jpayne@68 10
jpayne@68 11 from rich.progress import (
jpayne@68 12 BarColumn, Progress, ProgressColumn, Text, TimeElapsedColumn, TimeRemainingColumn, filesize)
jpayne@68 13
jpayne@68 14 from .std import TqdmExperimentalWarning
jpayne@68 15 from .std import tqdm as std_tqdm
jpayne@68 16
jpayne@68 17 __author__ = {"github.com/": ["casperdcl"]}
jpayne@68 18 __all__ = ['tqdm_rich', 'trrange', 'tqdm', 'trange']
jpayne@68 19
jpayne@68 20
jpayne@68 21 class FractionColumn(ProgressColumn):
jpayne@68 22 """Renders completed/total, e.g. '0.5/2.3 G'."""
jpayne@68 23 def __init__(self, unit_scale=False, unit_divisor=1000):
jpayne@68 24 self.unit_scale = unit_scale
jpayne@68 25 self.unit_divisor = unit_divisor
jpayne@68 26 super().__init__()
jpayne@68 27
jpayne@68 28 def render(self, task):
jpayne@68 29 """Calculate common unit for completed and total."""
jpayne@68 30 completed = int(task.completed)
jpayne@68 31 total = int(task.total)
jpayne@68 32 if self.unit_scale:
jpayne@68 33 unit, suffix = filesize.pick_unit_and_suffix(
jpayne@68 34 total,
jpayne@68 35 ["", "K", "M", "G", "T", "P", "E", "Z", "Y"],
jpayne@68 36 self.unit_divisor,
jpayne@68 37 )
jpayne@68 38 else:
jpayne@68 39 unit, suffix = filesize.pick_unit_and_suffix(total, [""], 1)
jpayne@68 40 precision = 0 if unit == 1 else 1
jpayne@68 41 return Text(
jpayne@68 42 f"{completed/unit:,.{precision}f}/{total/unit:,.{precision}f} {suffix}",
jpayne@68 43 style="progress.download")
jpayne@68 44
jpayne@68 45
jpayne@68 46 class RateColumn(ProgressColumn):
jpayne@68 47 """Renders human readable transfer speed."""
jpayne@68 48 def __init__(self, unit="", unit_scale=False, unit_divisor=1000):
jpayne@68 49 self.unit = unit
jpayne@68 50 self.unit_scale = unit_scale
jpayne@68 51 self.unit_divisor = unit_divisor
jpayne@68 52 super().__init__()
jpayne@68 53
jpayne@68 54 def render(self, task):
jpayne@68 55 """Show data transfer speed."""
jpayne@68 56 speed = task.speed
jpayne@68 57 if speed is None:
jpayne@68 58 return Text(f"? {self.unit}/s", style="progress.data.speed")
jpayne@68 59 if self.unit_scale:
jpayne@68 60 unit, suffix = filesize.pick_unit_and_suffix(
jpayne@68 61 speed,
jpayne@68 62 ["", "K", "M", "G", "T", "P", "E", "Z", "Y"],
jpayne@68 63 self.unit_divisor,
jpayne@68 64 )
jpayne@68 65 else:
jpayne@68 66 unit, suffix = filesize.pick_unit_and_suffix(speed, [""], 1)
jpayne@68 67 precision = 0 if unit == 1 else 1
jpayne@68 68 return Text(f"{speed/unit:,.{precision}f} {suffix}{self.unit}/s",
jpayne@68 69 style="progress.data.speed")
jpayne@68 70
jpayne@68 71
jpayne@68 72 class tqdm_rich(std_tqdm): # pragma: no cover
jpayne@68 73 """Experimental rich.progress GUI version of tqdm!"""
jpayne@68 74 # TODO: @classmethod: write()?
jpayne@68 75 def __init__(self, *args, **kwargs):
jpayne@68 76 """
jpayne@68 77 This class accepts the following parameters *in addition* to
jpayne@68 78 the parameters accepted by `tqdm`.
jpayne@68 79
jpayne@68 80 Parameters
jpayne@68 81 ----------
jpayne@68 82 progress : tuple, optional
jpayne@68 83 arguments for `rich.progress.Progress()`.
jpayne@68 84 options : dict, optional
jpayne@68 85 keyword arguments for `rich.progress.Progress()`.
jpayne@68 86 """
jpayne@68 87 kwargs = kwargs.copy()
jpayne@68 88 kwargs['gui'] = True
jpayne@68 89 # convert disable = None to False
jpayne@68 90 kwargs['disable'] = bool(kwargs.get('disable', False))
jpayne@68 91 progress = kwargs.pop('progress', None)
jpayne@68 92 options = kwargs.pop('options', {}).copy()
jpayne@68 93 super().__init__(*args, **kwargs)
jpayne@68 94
jpayne@68 95 if self.disable:
jpayne@68 96 return
jpayne@68 97
jpayne@68 98 warn("rich is experimental/alpha", TqdmExperimentalWarning, stacklevel=2)
jpayne@68 99 d = self.format_dict
jpayne@68 100 if progress is None:
jpayne@68 101 progress = (
jpayne@68 102 "[progress.description]{task.description}"
jpayne@68 103 "[progress.percentage]{task.percentage:>4.0f}%",
jpayne@68 104 BarColumn(bar_width=None),
jpayne@68 105 FractionColumn(
jpayne@68 106 unit_scale=d['unit_scale'], unit_divisor=d['unit_divisor']),
jpayne@68 107 "[", TimeElapsedColumn(), "<", TimeRemainingColumn(),
jpayne@68 108 ",", RateColumn(unit=d['unit'], unit_scale=d['unit_scale'],
jpayne@68 109 unit_divisor=d['unit_divisor']), "]"
jpayne@68 110 )
jpayne@68 111 options.setdefault('transient', not self.leave)
jpayne@68 112 self._prog = Progress(*progress, **options)
jpayne@68 113 self._prog.__enter__()
jpayne@68 114 self._task_id = self._prog.add_task(self.desc or "", **d)
jpayne@68 115
jpayne@68 116 def close(self):
jpayne@68 117 if self.disable:
jpayne@68 118 return
jpayne@68 119 self.display() # print 100%, vis #1306
jpayne@68 120 super().close()
jpayne@68 121 self._prog.__exit__(None, None, None)
jpayne@68 122
jpayne@68 123 def clear(self, *_, **__):
jpayne@68 124 pass
jpayne@68 125
jpayne@68 126 def display(self, *_, **__):
jpayne@68 127 if not hasattr(self, '_prog'):
jpayne@68 128 return
jpayne@68 129 self._prog.update(self._task_id, completed=self.n, description=self.desc)
jpayne@68 130
jpayne@68 131 def reset(self, total=None):
jpayne@68 132 """
jpayne@68 133 Resets to 0 iterations for repeated use.
jpayne@68 134
jpayne@68 135 Parameters
jpayne@68 136 ----------
jpayne@68 137 total : int or float, optional. Total to use for the new bar.
jpayne@68 138 """
jpayne@68 139 if hasattr(self, '_prog'):
jpayne@68 140 self._prog.reset(total=total)
jpayne@68 141 super().reset(total=total)
jpayne@68 142
jpayne@68 143
jpayne@68 144 def trrange(*args, **kwargs):
jpayne@68 145 """Shortcut for `tqdm.rich.tqdm(range(*args), **kwargs)`."""
jpayne@68 146 return tqdm_rich(range(*args), **kwargs)
jpayne@68 147
jpayne@68 148
jpayne@68 149 # Aliases
jpayne@68 150 tqdm = tqdm_rich
jpayne@68 151 trange = trrange