From daf36893eeef8a11a59023d85f86491fff882288 Mon Sep 17 00:00:00 2001 From: Casper da Costa-Luis Date: Sun, 10 Jan 2021 23:39:25 +0000 Subject: [PATCH] add `rich` --- .meta/.readme.rst | 3 + README.rst | 3 + environment.yml | 3 +- tests/tests_rich.py | 10 +++ tox.ini | 4 +- tqdm/rich.py | 157 ++++++++++++++++++++++++++++++++++++++++++++ 6 files changed, 178 insertions(+), 2 deletions(-) create mode 100644 tests/tests_rich.py create mode 100644 tqdm/rich.py diff --git a/.meta/.readme.rst b/.meta/.readme.rst index 29e942ba..8b1097b1 100644 --- a/.meta/.readme.rst +++ b/.meta/.readme.rst @@ -456,6 +456,9 @@ Submodules class tqdm.tk.tqdm(tqdm.tqdm): """Tkinter GUI version.""" + class tqdm.rich.tqdm(tqdm.tqdm): + """`rich.progress` version.""" + class tqdm.keras.TqdmCallback(keras.callbacks.Callback): """`keras` callback for epoch and batch progress.""" diff --git a/README.rst b/README.rst index c1d8440e..5e377468 100644 --- a/README.rst +++ b/README.rst @@ -673,6 +673,9 @@ Submodules class tqdm.tk.tqdm(tqdm.tqdm): """Tkinter GUI version.""" + class tqdm.rich.tqdm(tqdm.tqdm): + """`rich.progress` version.""" + class tqdm.keras.TqdmCallback(keras.callbacks.Callback): """`keras` callback for epoch and batch progress.""" diff --git a/environment.yml b/environment.yml index 2c74a3cf..a88f03f1 100644 --- a/environment.yml +++ b/environment.yml @@ -30,8 +30,9 @@ dependencies: - pandas - tensorflow # keras - requests # contrib.telegram -- twine # `pymake pypi` +- rich # rich - argopt # `cd wiki && pymake` +- twine # `pymake pypi` - wheel # `setup.py bdist_wheel` - pip: - py-make >=0.1.0 # `setup.py make/pymake` diff --git a/tests/tests_rich.py b/tests/tests_rich.py new file mode 100644 index 00000000..c75e246d --- /dev/null +++ b/tests/tests_rich.py @@ -0,0 +1,10 @@ +"""Test `tqdm.rich`.""" +import sys + +from .tests_tqdm import importorskip, mark + + +@mark.skipif(sys.version_info[:3] < (3, 6, 1), reason="`rich` needs py>=3.6.1") +def test_rich_import(): + """Test `tqdm.rich` import""" + importorskip('tqdm.rich') diff --git a/tox.ini b/tox.ini index d27c3a87..3adb9f01 100644 --- a/tox.ini +++ b/tox.ini @@ -33,10 +33,12 @@ passenv=CI TOXENV CODECOV_* COVERALLS_* CODACY_* HOME deps= {[extra]deps} cython + matplotlib numpy pandas - tf: tensorflow keras: keras + py{36,37,38,39}: rich + tf: tensorflow commands={[extra]commands} allowlist_externals={[extra]allowlist_externals} diff --git a/tqdm/rich.py b/tqdm/rich.py new file mode 100644 index 00000000..944691b9 --- /dev/null +++ b/tqdm/rich.py @@ -0,0 +1,157 @@ +""" +`rich.progress` decorator for iterators. + +Usage: +>>> from tqdm.rich import trange, tqdm +>>> for i in trange(10): +... ... +""" +from __future__ import absolute_import + +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 +from .utils import _range + +__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()`. + """ + kwargs = kwargs.copy() + kwargs['gui'] = True + # convert disable = None to False + kwargs['disable'] = bool(kwargs.get('disable', False)) + progress = kwargs.pop('progress', None) + 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']), "]" + ) + self._prog = Progress(*progress, transient=not self.leave) + self._prog.__enter__() + self._task_id = self._prog.add_task(self.desc or "", **d) + + def close(self, *args, **kwargs): + if self.disable: + return + super(tqdm_rich, self).close(*args, **kwargs) + 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): + """ + A shortcut for `tqdm.rich.tqdm(xrange(*args), **kwargs)`. + On Python3+, `range` is used instead of `xrange`. + """ + return tqdm_rich(_range(*args), **kwargs) + + +# Aliases +tqdm = tqdm_rich +trange = trrange