codekingpro/portable-devtools

1"""2`rich.progress` decorator for iterators.3 4Usage:5>>> from tqdm.rich import trange, tqdm6>>> for i in trange(10):7...     ...8"""9from warnings import warn10 11from rich.progress import (12    BarColumn, Progress, ProgressColumn, Text, TimeElapsedColumn, TimeRemainingColumn, filesize)13 14from .std import TqdmExperimentalWarning15from .std import tqdm as std_tqdm16 17__author__ = {"github.com/": ["casperdcl"]}18__all__ = ['tqdm_rich', 'trrange', 'tqdm', 'trange']19 20 21class FractionColumn(ProgressColumn):22    """Renders completed/total, e.g. '0.5/2.3 G'."""23    def __init__(self, unit_scale=False, unit_divisor=1000):24        self.unit_scale = unit_scale25        self.unit_divisor = unit_divisor26        super().__init__()27 28    def render(self, task):29        """Calculate common unit for completed and total."""30        completed = int(task.completed)31        total = int(task.total)32        if self.unit_scale:33            unit, suffix = filesize.pick_unit_and_suffix(34                total,35                ["", "K", "M", "G", "T", "P", "E", "Z", "Y"],36                self.unit_divisor,37            )38        else:39            unit, suffix = filesize.pick_unit_and_suffix(total, [""], 1)40        precision = 0 if unit == 1 else 141        return Text(42            f"{completed/unit:,.{precision}f}/{total/unit:,.{precision}f} {suffix}",43            style="progress.download")44 45 46class RateColumn(ProgressColumn):47    """Renders human readable transfer speed."""48    def __init__(self, unit="", unit_scale=False, unit_divisor=1000):49        self.unit = unit50        self.unit_scale = unit_scale51        self.unit_divisor = unit_divisor52        super().__init__()53 54    def render(self, task):55        """Show data transfer speed."""56        speed = task.speed57        if speed is None:58            return Text(f"? {self.unit}/s", style="progress.data.speed")59        if self.unit_scale:60            unit, suffix = filesize.pick_unit_and_suffix(61                speed,62                ["", "K", "M", "G", "T", "P", "E", "Z", "Y"],63                self.unit_divisor,64            )65        else:66            unit, suffix = filesize.pick_unit_and_suffix(speed, [""], 1)67        precision = 0 if unit == 1 else 168        return Text(f"{speed/unit:,.{precision}f} {suffix}{self.unit}/s",69                    style="progress.data.speed")70 71 72class tqdm_rich(std_tqdm):  # pragma: no cover73    """Experimental rich.progress GUI version of tqdm!"""74    # TODO: @classmethod: write()?75    def __init__(self, *args, **kwargs):76        """77        This class accepts the following parameters *in addition* to78        the parameters accepted by `tqdm`.79 80        Parameters81        ----------82        progress  : tuple, optional83            arguments for `rich.progress.Progress()`.84        options  : dict, optional85            keyword arguments for `rich.progress.Progress()`.86        """87        kwargs = kwargs.copy()88        kwargs['gui'] = True89        # convert disable = None to False90        kwargs['disable'] = bool(kwargs.get('disable', False))91        progress = kwargs.pop('progress', None)92        options = kwargs.pop('options', {}).copy()93        super().__init__(*args, **kwargs)94 95        if self.disable:96            return97 98        warn("rich is experimental/alpha", TqdmExperimentalWarning, stacklevel=2)99        d = self.format_dict100        if progress is None:101            progress = (102                "[progress.description]{task.description}"103                "[progress.percentage]{task.percentage:>4.0f}%",104                BarColumn(bar_width=None),105                FractionColumn(106                    unit_scale=d['unit_scale'], unit_divisor=d['unit_divisor']),107                "[", TimeElapsedColumn(), "<", TimeRemainingColumn(),108                ",", RateColumn(unit=d['unit'], unit_scale=d['unit_scale'],109                                unit_divisor=d['unit_divisor']), "]"110            )111        options.setdefault('transient', not self.leave)112        self._prog = Progress(*progress, **options)113        self._prog.__enter__()114        self._task_id = self._prog.add_task(self.desc or "", **d)115 116    def close(self):117        if self.disable:118            return119        self.display()  # print 100%, vis #1306120        super().close()121        self._prog.__exit__(None, None, None)122 123    def clear(self, *_, **__):124        pass125 126    def display(self, *_, **__):127        if not hasattr(self, '_prog'):128            return129        self._prog.update(self._task_id, completed=self.n, description=self.desc)130 131    def reset(self, total=None):132        """133        Resets to 0 iterations for repeated use.134 135        Parameters136        ----------137        total  : int or float, optional. Total to use for the new bar.138        """139        if hasattr(self, '_prog'):140            self._prog.reset(total=total)141        super().reset(total=total)142 143 144def trrange(*args, **kwargs):145    """Shortcut for `tqdm.rich.tqdm(range(*args), **kwargs)`."""146    return tqdm_rich(range(*args), **kwargs)147 148 149# Aliases150tqdm = tqdm_rich151trange = trrange152