from sys import stdout
from termcolor import colored as col
from framed_text.framed_text import frame_draw_top, frame_draw_bottom
from framed_text.shorten_text import ShortenText
from framed_text.utils import (
_validate_color,
ANSIEscape,
_format_col_text,
_remove_ansi,
_shorten_framed_title,
_ansi_pipe_clean,
_get_terminal_width,
_get_terminal_height,
)
[docs]
class ProgressBar:
[docs]
def __init__(self, value: int,
total: int,
label: str = '',
show_frame: bool = False,
show_percent: bool = True,
show_values: bool = False,
progress_icon: str = '#',
empty_icon: str = '.',
brackets: tuple[str, str] = ('[', ']'),
value_color: str | tuple[int, int, int] | None = None,
label_color: str | tuple[int, int, int] | None = None,
percent_color: str | tuple[int, int, int] | None = None,
progress_color: str | tuple[int, int, int] | None = None,
brackets_color: str | tuple[int, int, int] | None = None,
frame_color: str | tuple[int, int, int] | None = None,
safe_mode: bool = False,
skip_init: bool = False):
"""
Display a progress bar that shows the progress from a value to a total.
The progress bar will snap to the bottom of the terminal.
:param value: Current value
:param total: Total value
:param label: Optional label for the Progress Bar
:param show_frame: If true, will show a frame around the progress bar, similar to ``FramedText``
:param show_percent: If true, will show percentage of progress. Will replace values display unless ``show_values`` is true
:param show_values: If true, will show current value and total value. If false and ``show_percent`` is false, will not display any values
:param progress_icon: Icon used for the progress display
:param empty_icon: Icon used for the empty space in the progress display
:param brackets: Brackets to surround the progress display. First value is the left bracket, second value is the right bracket
:param value_color: Color of the label
:param label_color: Color of the label
:param percent_color: Color of the percentage
:param progress_color: Color of the progress display
:param brackets_color: Color of the brackets
:param frame_color: Color of the frame
:param safe_mode: If true, will clear the terminal and move cursor to top to ensure room for progress bar and any other data.
:param skip_init: If true, will only initialize when ``display`` or ``update_progress`` is called
"""
self._value: int = value
self._total: int = total
self._label: str = label
self._show_frame: bool = show_frame
self._show_percent: bool = show_percent
self._show_values: bool = show_values
self._progress_icon: str = progress_icon
self._empty_icon: str = empty_icon
self._brackets: tuple[str, str] = brackets
self._value_color: str | tuple[int, int, int] | None = None
self._label_color: str | tuple[int, int, int] | None = None
self._percent_color: str | tuple[int, int, int] | None = None
self._progress_color: str | tuple[int, int, int] | None = None
self._brackets_color: str | tuple[int, int, int] | None = None
self._frame_color: str | tuple[int, int, int] | None = None
self._safe_mode: bool = safe_mode
self._skip_init: bool = skip_init
self._prog_size: int = 1
self._limit: int = _get_terminal_width()
self._term_is_init: bool = False
self._term_lines: int = 0
self._term_columns: int = 0
self._term_columns: int = _get_terminal_width()
self._term_lines: int = _get_terminal_height()
self.bar_display: list[str] = []
# Validate colors
if value_color:
_validate_color(color=value_color)
self._value_color = value_color.lower() if isinstance(value_color, str) else value_color
if label_color:
_validate_color(color=label_color)
self._label_color = label_color.lower() if isinstance(label_color, str) else label_color
if percent_color:
_validate_color(color=percent_color)
self._percent_color = percent_color.lower() if isinstance(percent_color, str) else percent_color
if progress_color:
_validate_color(color=progress_color)
self._progress_color = progress_color.lower() if isinstance(progress_color, str) else progress_color
if brackets_color:
_validate_color(color=brackets_color)
self._brackets_color = brackets_color.lower() if isinstance(brackets_color, str) else brackets_color
if frame_color:
_validate_color(color=frame_color)
self._frame_color = frame_color.lower() if isinstance(frame_color, str) else frame_color
# Setup progress bar size
if self._label and not self._show_frame:
self._prog_size += 1
if self._show_frame:
self._prog_size += 2
self._limit -= 4
# Initialize terminal
if not self._term_is_init and not self._skip_init:
self._init_term()
[docs]
def update_progress(self, value: int | None = None, total: int | None = None):
# some code
if value:
self._value = value
if total:
self._total = total
self.display()
[docs]
def display(self):
"""
Display the progress bar
"""
if not self._term_is_init:
self._init_term()
self._term_is_init = True
self._display_progress()
[docs]
def reset(self):
# Resets the terminal to its original state
self._term_is_init = False
# 1. Save current position
stdout.write(ANSIEscape.SAVE_POS)
# 2. Reset scrollable region
stdout.write(ANSIEscape.reset_scrollable_region())
# 3. Move cursor to affected lines
for line_num in range(self._prog_size):
stdout.write(ANSIEscape.move_to(line=self._term_lines - line_num))
# 4. Clear affected lines
stdout.write(ANSIEscape.ERASE_LINE)
# 5. Restore cursor position
stdout.write(ANSIEscape.RESTORE_POS)
def _init_term(self):
# Initializes the terminal in a state for the progress bar
self._term_is_init = True
if self._safe_mode:
# Clear screen and move to top
stdout.write(ANSIEscape.ERASE_SCREEN)
stdout.write(ANSIEscape.move_to(line=0, column=0))
else:
# Create new empty lines, then move back up.
# This ensures there is enough room for any data outside the progress bar
# to be printed correctly as now the cursor will be
# within the scrollable region when it is created.
stdout.write('\n' * self._prog_size)
stdout.write(ANSIEscape.move_up(self._prog_size))
# 1. Create space for the progress bar using scrollbar
stdout.write('\n')
# 2. Save current position
stdout.write(ANSIEscape.SAVE_POS)
# 3. Set scrollable region
stdout.write(ANSIEscape.set_scrollable_region(top=0, bottom=self._term_lines - self._prog_size))
# 4. Restore cursor position
stdout.write(ANSIEscape.RESTORE_POS)
# 5. Move cursor up
stdout.write(ANSIEscape.move_up())
def _display_progress(self):
# Internal. Just displays the progress bar
# Calculate available space for the progress bar
# Clear display
self.bar_display.clear()
# Format Label
if self._label and not self._show_frame:
self._label = ShortenText(text=_remove_ansi(text=self._label), limit=self._limit).__str__()
elif self._label:
self._label = _shorten_framed_title(title=self._label, limit=self._limit - 2)
label = _format_col_text(text=self._label, color=self._label_color)
# Format values or percent display
value: str = ''
percent: int = round((self._value / self._total) * 100)
if self._show_values and self._show_percent:
# Both
value = f" {self._value}/{self._total} {' ' * (3 - len(str(percent)))}({percent}%)"
elif self._show_values:
# Values
value = f" {self._value}/{self._total}"
elif self._show_percent:
# Percent
value = f" {' ' * (3 - len(str(percent)))}{percent}%"
# Format progress bar. Ensure only 1 char is used
space_available: int = self._limit - len(value)
space_available -= 2 # Account for brackets
prog_count: int = round((self._value / self._total) * space_available)
filled: str = self._progress_icon[0] * prog_count
empty: str = self._empty_icon[0] * (space_available - prog_count)
bar: str = f"{self._brackets[0][0]}{filled}{empty}{self._brackets[1][0]}{value}"
# Combine
_line: str = ''
if self._show_frame:
if self._label:
# Label becomes frame title
self.bar_display.append(
frame_draw_top(title=label, title_len=len(label), frame_color=self._frame_color))
else:
self.bar_display.append(frame_draw_top(frame_color=self._frame_color))
_line = f"{col('│', self._frame_color)} {bar} {col('│', self._frame_color)}"
self.bar_display.append(_ansi_pipe_clean(text=_line))
self.bar_display.append(frame_draw_bottom(frame_color=self._frame_color))
else:
if self._label:
self.bar_display.append(f"{label}")
self.bar_display.append(f"{bar}")
# Save cursor position
stdout.write(ANSIEscape.SAVE_POS)
for i, line in enumerate(self.bar_display):
# Move to line
stdout.write(ANSIEscape.move_to(line=self._term_lines - (self._prog_size - i - 1)))
# Clear line
stdout.write(ANSIEscape.ERASE_LINE)
# Print line
stdout.write(line)
stdout.write(ANSIEscape.RESTORE_POS)
# Restore cursor position
stdout.flush()
stdout.write(ANSIEscape.RESTORE_POS)
def __str__(self) -> str:
return '\n'.join(self.bar_display)
[docs]
@staticmethod
def force_reset(show_frame: bool = False, label: bool = False):
"""
Force resets the terminal. This should only be used in the ``except`` block of a try-except block.
Make sure to pass the same values to this function that were passed to the ``ProgressBar`` constructor,
otherwise the reset may not work as expected.
**NOTE**: This method does not change ``ProgressBar._term_is_init``. If you plan to re-use the progress bar
after a force reset, you will need to call ``ProgressBar._init_term`` manually.
This function behaves like ``ProgressBar.reset``
:param show_frame: Whether ``show_frame`` was enabled for the progress bar.
:param label: Whether ``label`` was enabled for the progress bar.
"""
prog_size: int = 1
prog_size += 1 if label and not show_frame else 0
prog_size += 2 if show_frame else 0
term_lines: int = _get_terminal_height()
# 1. Save current position
stdout.write(ANSIEscape.SAVE_POS)
# 2. Reset scrollable region
stdout.write(ANSIEscape.reset_scrollable_region())
# 3. Move cursor to affected lines
for line_num in range(prog_size):
stdout.write(ANSIEscape.move_to(line=term_lines - line_num))
# 4. Clear affected lines
stdout.write(ANSIEscape.ERASE_LINE)
# 5. Restore cursor position
stdout.write(ANSIEscape.RESTORE_POS)