"""
Formatter classes for the progress bar.
Each progress bar consists of a list of these formatters.
"""
import datetime
import time
from abc import ABCMeta, abstractmethod
from typing import TYPE_CHECKING, List, Tuple
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:>5}%"
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 = "{start}{bar_a}{bar_b}{bar_c}{end}"
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:>3}/{total:>3}"
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}").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}"
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:.2f}"
)
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("{0}").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(" "),
]