jpayne@68: """ jpayne@68: `rich.progress` decorator for iterators. jpayne@68: jpayne@68: Usage: jpayne@68: >>> from tqdm.rich import trange, tqdm jpayne@68: >>> for i in trange(10): jpayne@68: ... ... jpayne@68: """ jpayne@68: from warnings import warn jpayne@68: jpayne@68: from rich.progress import ( jpayne@68: BarColumn, Progress, ProgressColumn, Text, TimeElapsedColumn, TimeRemainingColumn, filesize) jpayne@68: jpayne@68: from .std import TqdmExperimentalWarning jpayne@68: from .std import tqdm as std_tqdm jpayne@68: jpayne@68: __author__ = {"github.com/": ["casperdcl"]} jpayne@68: __all__ = ['tqdm_rich', 'trrange', 'tqdm', 'trange'] jpayne@68: jpayne@68: jpayne@68: class FractionColumn(ProgressColumn): jpayne@68: """Renders completed/total, e.g. '0.5/2.3 G'.""" jpayne@68: def __init__(self, unit_scale=False, unit_divisor=1000): jpayne@68: self.unit_scale = unit_scale jpayne@68: self.unit_divisor = unit_divisor jpayne@68: super().__init__() jpayne@68: jpayne@68: def render(self, task): jpayne@68: """Calculate common unit for completed and total.""" jpayne@68: completed = int(task.completed) jpayne@68: total = int(task.total) jpayne@68: if self.unit_scale: jpayne@68: unit, suffix = filesize.pick_unit_and_suffix( jpayne@68: total, jpayne@68: ["", "K", "M", "G", "T", "P", "E", "Z", "Y"], jpayne@68: self.unit_divisor, jpayne@68: ) jpayne@68: else: jpayne@68: unit, suffix = filesize.pick_unit_and_suffix(total, [""], 1) jpayne@68: precision = 0 if unit == 1 else 1 jpayne@68: return Text( jpayne@68: f"{completed/unit:,.{precision}f}/{total/unit:,.{precision}f} {suffix}", jpayne@68: style="progress.download") jpayne@68: jpayne@68: jpayne@68: class RateColumn(ProgressColumn): jpayne@68: """Renders human readable transfer speed.""" jpayne@68: def __init__(self, unit="", unit_scale=False, unit_divisor=1000): jpayne@68: self.unit = unit jpayne@68: self.unit_scale = unit_scale jpayne@68: self.unit_divisor = unit_divisor jpayne@68: super().__init__() jpayne@68: jpayne@68: def render(self, task): jpayne@68: """Show data transfer speed.""" jpayne@68: speed = task.speed jpayne@68: if speed is None: jpayne@68: return Text(f"? {self.unit}/s", style="progress.data.speed") jpayne@68: if self.unit_scale: jpayne@68: unit, suffix = filesize.pick_unit_and_suffix( jpayne@68: speed, jpayne@68: ["", "K", "M", "G", "T", "P", "E", "Z", "Y"], jpayne@68: self.unit_divisor, jpayne@68: ) jpayne@68: else: jpayne@68: unit, suffix = filesize.pick_unit_and_suffix(speed, [""], 1) jpayne@68: precision = 0 if unit == 1 else 1 jpayne@68: return Text(f"{speed/unit:,.{precision}f} {suffix}{self.unit}/s", jpayne@68: style="progress.data.speed") jpayne@68: jpayne@68: jpayne@68: class tqdm_rich(std_tqdm): # pragma: no cover jpayne@68: """Experimental rich.progress GUI version of tqdm!""" jpayne@68: # TODO: @classmethod: write()? jpayne@68: def __init__(self, *args, **kwargs): jpayne@68: """ jpayne@68: This class accepts the following parameters *in addition* to jpayne@68: the parameters accepted by `tqdm`. jpayne@68: jpayne@68: Parameters jpayne@68: ---------- jpayne@68: progress : tuple, optional jpayne@68: arguments for `rich.progress.Progress()`. jpayne@68: options : dict, optional jpayne@68: keyword arguments for `rich.progress.Progress()`. jpayne@68: """ jpayne@68: kwargs = kwargs.copy() jpayne@68: kwargs['gui'] = True jpayne@68: # convert disable = None to False jpayne@68: kwargs['disable'] = bool(kwargs.get('disable', False)) jpayne@68: progress = kwargs.pop('progress', None) jpayne@68: options = kwargs.pop('options', {}).copy() jpayne@68: super().__init__(*args, **kwargs) jpayne@68: jpayne@68: if self.disable: jpayne@68: return jpayne@68: jpayne@68: warn("rich is experimental/alpha", TqdmExperimentalWarning, stacklevel=2) jpayne@68: d = self.format_dict jpayne@68: if progress is None: jpayne@68: progress = ( jpayne@68: "[progress.description]{task.description}" jpayne@68: "[progress.percentage]{task.percentage:>4.0f}%", jpayne@68: BarColumn(bar_width=None), jpayne@68: FractionColumn( jpayne@68: unit_scale=d['unit_scale'], unit_divisor=d['unit_divisor']), jpayne@68: "[", TimeElapsedColumn(), "<", TimeRemainingColumn(), jpayne@68: ",", RateColumn(unit=d['unit'], unit_scale=d['unit_scale'], jpayne@68: unit_divisor=d['unit_divisor']), "]" jpayne@68: ) jpayne@68: options.setdefault('transient', not self.leave) jpayne@68: self._prog = Progress(*progress, **options) jpayne@68: self._prog.__enter__() jpayne@68: self._task_id = self._prog.add_task(self.desc or "", **d) jpayne@68: jpayne@68: def close(self): jpayne@68: if self.disable: jpayne@68: return jpayne@68: self.display() # print 100%, vis #1306 jpayne@68: super().close() jpayne@68: self._prog.__exit__(None, None, None) jpayne@68: jpayne@68: def clear(self, *_, **__): jpayne@68: pass jpayne@68: jpayne@68: def display(self, *_, **__): jpayne@68: if not hasattr(self, '_prog'): jpayne@68: return jpayne@68: self._prog.update(self._task_id, completed=self.n, description=self.desc) jpayne@68: jpayne@68: def reset(self, total=None): jpayne@68: """ jpayne@68: Resets to 0 iterations for repeated use. jpayne@68: jpayne@68: Parameters jpayne@68: ---------- jpayne@68: total : int or float, optional. Total to use for the new bar. jpayne@68: """ jpayne@68: if hasattr(self, '_prog'): jpayne@68: self._prog.reset(total=total) jpayne@68: super().reset(total=total) jpayne@68: jpayne@68: jpayne@68: def trrange(*args, **kwargs): jpayne@68: """Shortcut for `tqdm.rich.tqdm(range(*args), **kwargs)`.""" jpayne@68: return tqdm_rich(range(*args), **kwargs) jpayne@68: jpayne@68: jpayne@68: # Aliases jpayne@68: tqdm = tqdm_rich jpayne@68: trange = trrange