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