1 files changed, 429 insertions, 0 deletions
diff --git a/src/prompt_toolkit/shortcuts/progress_bar/formatters.py b/src/prompt_toolkit/shortcuts/progress_bar/formatters.py
new file mode 100644
index 0000000..dd0339c
--- /dev/null
+++ b/src/prompt_toolkit/shortcuts/progress_bar/formatters.py
@@ -0,0 +1,429 @@
+"""
+Formatter classes for the progress bar.
+Each progress bar consists of a list of these formatters.
+"""
+from __future__ import annotations
+
+import datetime
+import time
+from abc import ABCMeta, abstractmethod
+from typing import TYPE_CHECKING
+
+from prompt_toolkit.formatted_text import (
+    HTML,
+    AnyFormattedText,
+    StyleAndTextTuples,
+    to_formatted_text,
+)
+from prompt_toolkit.formatted_text.utils import fragment_list_width
+from prompt_toolkit.layout.dimension import AnyDimension, D
+from prompt_toolkit.layout.utils import explode_text_fragments
+from prompt_toolkit.utils import get_cwidth
+
+if TYPE_CHECKING:
+    from .base import ProgressBar, ProgressBarCounter
+
+__all__ = [
+    "Formatter",
+    "Text",
+    "Label",
+    "Percentage",
+    "Bar",
+    "Progress",
+    "TimeElapsed",
+    "TimeLeft",
+    "IterationsPerSecond",
+    "SpinningWheel",
+    "Rainbow",
+    "create_default_formatters",
+]
+
+
+class Formatter(metaclass=ABCMeta):
+    """
+    Base class for any formatter.
+    """
+
+    @abstractmethod
+    def format(
+        self,
+        progress_bar: ProgressBar,
+        progress: ProgressBarCounter[object],
+        width: int,
+    ) -> AnyFormattedText:
+        pass
+
+    def get_width(self, progress_bar: ProgressBar) -> AnyDimension:
+        return D()
+
+
+class Text(Formatter):
+    """
+    Display plain text.
+    """
+
+    def __init__(self, text: AnyFormattedText, style: str = "") -> None:
+        self.text = to_formatted_text(text, style=style)
+
+    def format(
+        self,
+        progress_bar: ProgressBar,
+        progress: ProgressBarCounter[object],
+        width: int,
+    ) -> AnyFormattedText:
+        return self.text
+
+    def get_width(self, progress_bar: ProgressBar) -> AnyDimension:
+        return fragment_list_width(self.text)
+
+
+class Label(Formatter):
+    """
+    Display the name of the current task.
+
+    :param width: If a `width` is given, use this width. Scroll the text if it
+        doesn't fit in this width.
+    :param suffix: String suffix to be added after the task name, e.g. ': '.
+        If no task name was given, no suffix will be added.
+    """
+
+    def __init__(self, width: AnyDimension = None, suffix: str = "") -> None:
+        self.width = width
+        self.suffix = suffix
+
+    def _add_suffix(self, label: AnyFormattedText) -> StyleAndTextTuples:
+        label = to_formatted_text(label, style="class:label")
+        return label + [("", self.suffix)]
+
+    def format(
+        self,
+        progress_bar: ProgressBar,
+        progress: ProgressBarCounter[object],
+        width: int,
+    ) -> AnyFormattedText:
+        label = self._add_suffix(progress.label)
+        cwidth = fragment_list_width(label)
+
+        if cwidth > width:
+            # It doesn't fit -> scroll task name.
+            label = explode_text_fragments(label)
+            max_scroll = cwidth - width
+            current_scroll = int(time.time() * 3 % max_scroll)
+            label = label[current_scroll:]
+
+        return label
+
+    def get_width(self, progress_bar: ProgressBar) -> AnyDimension:
+        if self.width:
+            return self.width
+
+        all_labels = [self._add_suffix(c.label) for c in progress_bar.counters]
+        if all_labels:
+            max_widths = max(fragment_list_width(l) for l in all_labels)
+            return D(preferred=max_widths, max=max_widths)
+        else:
+            return D()
+
+
+class Percentage(Formatter):
+    """
+    Display the progress as a percentage.
+    """
+
+    template = "<percentage>{percentage:>5}%</percentage>"
+
+    def format(
+        self,
+        progress_bar: ProgressBar,
+        progress: ProgressBarCounter[object],
+        width: int,
+    ) -> AnyFormattedText:
+        return HTML(self.template).format(percentage=round(progress.percentage, 1))
+
+    def get_width(self, progress_bar: ProgressBar) -> AnyDimension:
+        return D.exact(6)
+
+
+class Bar(Formatter):
+    """
+    Display the progress bar itself.
+    """
+
+    template = "<bar>{start}<bar-a>{bar_a}</bar-a><bar-b>{bar_b}</bar-b><bar-c>{bar_c}</bar-c>{end}</bar>"
+
+    def __init__(
+        self,
+        start: str = "[",
+        end: str = "]",
+        sym_a: str = "=",
+        sym_b: str = ">",
+        sym_c: str = " ",
+        unknown: str = "#",
+    ) -> None:
+        assert len(sym_a) == 1 and get_cwidth(sym_a) == 1
+        assert len(sym_c) == 1 and get_cwidth(sym_c) == 1
+
+        self.start = start
+        self.end = end
+        self.sym_a = sym_a
+        self.sym_b = sym_b
+        self.sym_c = sym_c
+        self.unknown = unknown
+
+    def format(
+        self,
+        progress_bar: ProgressBar,
+        progress: ProgressBarCounter[object],
+        width: int,
+    ) -> AnyFormattedText:
+        if progress.done or progress.total or progress.stopped:
+            sym_a, sym_b, sym_c = self.sym_a, self.sym_b, self.sym_c
+
+            # Compute pb_a based on done, total, or stopped states.
+            if progress.done:
+                # 100% completed irrelevant of how much was actually marked as completed.
+                percent = 1.0
+            else:
+                # Show percentage completed.
+                percent = progress.percentage / 100
+        else:
+            # Total is unknown and bar is still running.
+            sym_a, sym_b, sym_c = self.sym_c, self.unknown, self.sym_c
+
+            # Compute percent based on the time.
+            percent = time.time() * 20 % 100 / 100
+
+        # Subtract left, sym_b, and right.
+        width -= get_cwidth(self.start + sym_b + self.end)
+
+        # Scale percent by width
+        pb_a = int(percent * width)
+        bar_a = sym_a * pb_a
+        bar_b = sym_b
+        bar_c = sym_c * (width - pb_a)
+
+        return HTML(self.template).format(
+            start=self.start, end=self.end, bar_a=bar_a, bar_b=bar_b, bar_c=bar_c
+        )
+
+    def get_width(self, progress_bar: ProgressBar) -> AnyDimension:
+        return D(min=9)
+
+
+class Progress(Formatter):
+    """
+    Display the progress as text.  E.g. "8/20"
+    """
+
+    template = "<current>{current:>3}</current>/<total>{total:>3}</total>"
+
+    def format(
+        self,
+        progress_bar: ProgressBar,
+        progress: ProgressBarCounter[object],
+        width: int,
+    ) -> AnyFormattedText:
+        return HTML(self.template).format(
+            current=progress.items_completed, total=progress.total or "?"
+        )
+
+    def get_width(self, progress_bar: ProgressBar) -> AnyDimension:
+        all_lengths = [
+            len("{:>3}".format(c.total or "?")) for c in progress_bar.counters
+        ]
+        all_lengths.append(1)
+        return D.exact(max(all_lengths) * 2 + 1)
+
+
+def _format_timedelta(timedelta: datetime.timedelta) -> str:
+    """
+    Return hh:mm:ss, or mm:ss if the amount of hours is zero.
+    """
+    result = f"{timedelta}".split(".")[0]
+    if result.startswith("0:"):
+        result = result[2:]
+    return result
+
+
+class TimeElapsed(Formatter):
+    """
+    Display the elapsed time.
+    """
+
+    def format(
+        self,
+        progress_bar: ProgressBar,
+        progress: ProgressBarCounter[object],
+        width: int,
+    ) -> AnyFormattedText:
+        text = _format_timedelta(progress.time_elapsed).rjust(width)
+        return HTML("<time-elapsed>{time_elapsed}</time-elapsed>").format(
+            time_elapsed=text
+        )
+
+    def get_width(self, progress_bar: ProgressBar) -> AnyDimension:
+        all_values = [
+            len(_format_timedelta(c.time_elapsed)) for c in progress_bar.counters
+        ]
+        if all_values:
+            return max(all_values)
+        return 0
+
+
+class TimeLeft(Formatter):
+    """
+    Display the time left.
+    """
+
+    template = "<time-left>{time_left}</time-left>"
+    unknown = "?:??:??"
+
+    def format(
+        self,
+        progress_bar: ProgressBar,
+        progress: ProgressBarCounter[object],
+        width: int,
+    ) -> AnyFormattedText:
+        time_left = progress.time_left
+        if time_left is not None:
+            formatted_time_left = _format_timedelta(time_left)
+        else:
+            formatted_time_left = self.unknown
+
+        return HTML(self.template).format(time_left=formatted_time_left.rjust(width))
+
+    def get_width(self, progress_bar: ProgressBar) -> AnyDimension:
+        all_values = [
+            len(_format_timedelta(c.time_left)) if c.time_left is not None else 7
+            for c in progress_bar.counters
+        ]
+        if all_values:
+            return max(all_values)
+        return 0
+
+
+class IterationsPerSecond(Formatter):
+    """
+    Display the iterations per second.
+    """
+
+    template = (
+        "<iterations-per-second>{iterations_per_second:.2f}</iterations-per-second>"
+    )
+
+    def format(
+        self,
+        progress_bar: ProgressBar,
+        progress: ProgressBarCounter[object],
+        width: int,
+    ) -> AnyFormattedText:
+        value = progress.items_completed / progress.time_elapsed.total_seconds()
+        return HTML(self.template.format(iterations_per_second=value))
+
+    def get_width(self, progress_bar: ProgressBar) -> AnyDimension:
+        all_values = [
+            len(f"{c.items_completed / c.time_elapsed.total_seconds():.2f}")
+            for c in progress_bar.counters
+        ]
+        if all_values:
+            return max(all_values)
+        return 0
+
+
+class SpinningWheel(Formatter):
+    """
+    Display a spinning wheel.
+    """
+
+    characters = r"/-\|"
+
+    def format(
+        self,
+        progress_bar: ProgressBar,
+        progress: ProgressBarCounter[object],
+        width: int,
+    ) -> AnyFormattedText:
+        index = int(time.time() * 3) % len(self.characters)
+        return HTML("<spinning-wheel>{0}</spinning-wheel>").format(
+            self.characters[index]
+        )
+
+    def get_width(self, progress_bar: ProgressBar) -> AnyDimension:
+        return D.exact(1)
+
+
+def _hue_to_rgb(hue: float) -> tuple[int, int, int]:
+    """
+    Take hue between 0 and 1, return (r, g, b).
+    """
+    i = int(hue * 6.0)
+    f = (hue * 6.0) - i
+
+    q = int(255 * (1.0 - f))
+    t = int(255 * (1.0 - (1.0 - f)))
+
+    i %= 6
+
+    return [
+        (255, t, 0),
+        (q, 255, 0),
+        (0, 255, t),
+        (0, q, 255),
+        (t, 0, 255),
+        (255, 0, q),
+    ][i]
+
+
+class Rainbow(Formatter):
+    """
+    For the fun. Add rainbow colors to any of the other formatters.
+    """
+
+    colors = ["#%.2x%.2x%.2x" % _hue_to_rgb(h / 100.0) for h in range(0, 100)]
+
+    def __init__(self, formatter: Formatter) -> None:
+        self.formatter = formatter
+
+    def format(
+        self,
+        progress_bar: ProgressBar,
+        progress: ProgressBarCounter[object],
+        width: int,
+    ) -> AnyFormattedText:
+        # Get formatted text from nested formatter, and explode it in
+        # text/style tuples.
+        result = self.formatter.format(progress_bar, progress, width)
+        result = explode_text_fragments(to_formatted_text(result))
+
+        # Insert colors.
+        result2: StyleAndTextTuples = []
+        shift = int(time.time() * 3) % len(self.colors)
+
+        for i, (style, text, *_) in enumerate(result):
+            result2.append(
+                (style + " " + self.colors[(i + shift) % len(self.colors)], text)
+            )
+        return result2
+
+    def get_width(self, progress_bar: ProgressBar) -> AnyDimension:
+        return self.formatter.get_width(progress_bar)
+
+
+def create_default_formatters() -> list[Formatter]:
+    """
+    Return the list of default formatters.
+    """
+    return [
+        Label(),
+        Text(" "),
+        Percentage(),
+        Text(" "),
+        Bar(),
+        Text(" "),
+        Progress(),
+        Text(" "),
+        Text("eta [", style="class:time-left"),
+        TimeLeft(),
+        Text("]", style="class:time-left"),
+        Text(" "),
+    ]