1 files changed, 156 insertions, 0 deletions
diff --git a/tqdm/rich.py b/tqdm/rich.py
new file mode 100644
index 0000000..69893ff
--- /dev/null
+++ b/tqdm/rich.py
@@ -0,0 +1,156 @@
+"""
+`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()`.
+        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):
+    """
+    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