Source code for framed_text.progress_bar

from sys import stdout

from termcolor import colored as col

from framed_text.color import Color
from framed_text.framed_text import _frame_draw_top, _frame_draw_bottom
from framed_text.shorten_text import ShortenText
from framed_text.utils import (
    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, safe_mode: bool = False, skip_init: bool = False, progress_icon: str = '#', empty_icon: str = '.', brackets: tuple[str, str] = ('[', ']'), value_color: str | tuple[int, int, int] | Color | None = None, label_color: str | tuple[int, int, int] | Color | None = None, percent_color: str | tuple[int, int, int] | Color | None = None, progress_color: str | tuple[int, int, int] | Color | None = None, empty_color: str | tuple[int, int, int] | Color | None = None, brackets_color: str | tuple[int, int, int] | Color | None = None, frame_color: str | tuple[int, int, int] | Color | None = None): """ 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 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. :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 empty_color: Color of the empty space in the progress display. :param brackets_color: Color of the brackets. :param frame_color: Color of the frame. """ 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._safe_mode: bool = safe_mode self._skip_init: bool = skip_init self._progress_icon: str = progress_icon self._empty_icon: str = empty_icon self._brackets: tuple[str, str] = brackets self._value_color: Color = Color(value_color) if value_color else Color() self._label_color: Color = Color(label_color) if label_color else Color() self._percent_color: Color = Color(percent_color) if percent_color else Color() self._progress_color: Color = Color(progress_color) if progress_color else Color() self._empty_color: Color = Color(empty_color) if empty_color else Color() self._brackets_color: Color = Color(brackets_color) if brackets_color else Color() self._frame_color: Color = Color(frame_color) if frame_color else Color() self._prog_size: int = 1 self._term_is_init: bool = False self._term_columns: int = _get_terminal_width() self._term_lines: int = _get_terminal_height() self._limit: int = self._term_columns self.bar_display: list[str] = [] # 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 # Shorten label if necessary label_no_ansi: str = _remove_ansi(text=self._label) if len(label_no_ansi) > self._limit: if self._show_frame: self._label = _shorten_framed_title(title=label_no_ansi, limit=self._limit - 2) else: self._label = ShortenText(text=label_no_ansi, limit=self._limit).__str__() # Initialize terminal if not self._term_is_init and not self._skip_init: self._init_term()
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 Data _label: str = _format_col_text(text=self._label, color=self._label_color.tc_color) _cur_val: str = _format_col_text(text=str(self._value), color=self._value_color.tc_color) _total_val: str = _format_col_text(text=str(self._total), color=self._value_color.tc_color) _percent_val: int = round((self._value / self._total) * 100) _percent: str = _format_col_text(text=f"{_percent_val}%", color=self._percent_color.tc_color) # Format values or percent display val_display: str = '' if self._show_values and self._show_percent: # Both val_display = f" {_cur_val}/{_total_val} {' ' * (3 - len(str(_percent_val)))}({_percent})" elif self._show_values: # Values val_display = f" {_cur_val}/{_total_val}" elif self._show_percent: # Percent val_display = f" {' ' * (3 - len(str(_percent_val)))}{_percent}" # Format progress bar. Ensure only 1 char is used _space_available: int = self._limit - len(_remove_ansi(text=val_display)) _space_available -= 2 # Account for brackets _prog_count: int = round((self._value / self._total) * _space_available) _filled: str = _format_col_text(text=self._progress_icon[0] * _prog_count, color=self._progress_color.tc_color) _empty: str = _format_col_text(text=self._empty_icon[0] * (_space_available - _prog_count), color=self._empty_color.tc_color) _bracket_s: str = _format_col_text(text=self._brackets[0][0], color=self._brackets_color.tc_color) _bracket_e: str = _format_col_text(text=self._brackets[1][0], color=self._brackets_color.tc_color) bar: str = f"{_bracket_s}{_filled}{_empty}{_bracket_e}{val_display}" # 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(_remove_ansi(text=_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.tc_color)} {bar} {col('│', self._frame_color.tc_color)}" self.bar_display.append(_ansi_pipe_clean(text=_line)) self.bar_display.append(_frame_draw_bottom(frame_color=self._frame_color)) else: if _label: self.bar_display.append(_label) self.bar_display.append(bar) # Save cursor position stdout.write(ANSIEscape.SAVE_POS) # Draw display 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] def update_progress(self, value: int | None = None, total: int | None = None): """ Update the values of the progress bar, then displays it :param value: New Value. Leave blank to use current value :param total: New Total. Leave blank to use current total """ 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: # Initialize terminal 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)
[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)