from multiprocessing import Process, Event
from sys import stdout
from time import sleep
from typing import Literal
from framed_text.color import Color
from framed_text.labeled_data import AnyLabeledData
from framed_text.shorten_text import ShortenText
from framed_text.status import AnyStatus
from framed_text.utils import (
ANSIEscape,
_format_col_text,
_get_terminal_width,
_get_terminal_height,
_remove_ansi
)
[docs]
class Spinner:
BUILTIN_ICONS: dict[str, list[str]] = {
"default": ['/', '-', '\\', '|'],
"cube": ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'],
"dots": [' ', '. ', '.. ', '...'],
"dots2": [' ', '. ', '.. ', '...', ' ..', ' .'],
}
[docs]
def __init__(self, text: str | AnyStatus | AnyLabeledData = '',
icons_id: str = "default",
done_icon: str = '',
custom_icons: list[str] | None = None,
interval: float = 0.2,
position: Literal["start", "end"] = "start",
indent: int = 1,
spinner_color: str | tuple[int, int, int] | Color | None = "yellow",
done_color: str | tuple[int, int, int] | Color | None = "green"):
"""
Draws a spinner on a screen
**Position**
- ``"start"``: The spinner will be placed at the start of the line.
- ``"end"``: The spinner will be placed at the end of the line.
:param text: Text to display next to the spinner. Omit for no text.
:param icons_id: ID string of one of the built-in icon sets from ``Spinner.BUILTIN_ICONS``
:param done_icon: Icon to use for when the spinner is complete
:param custom_icons: A list of strings to use as the spinner icon. Overrides ``icons_id``
:param interval: The delay between spinner icon updates.
:param position: Where the spinner should be placed on a line.
:param indent: How many spaces the line should be indented.
:param spinner_color: Color of the spinner icon while active.
:param done_color: Color for the spinner icon when stopped.
"""
self._text: str | AnyStatus | AnyLabeledData = text
self._icons_id: str = icons_id
self._done_icon: str = _remove_ansi(text=done_icon.__str__())
self._icons: list[str] = custom_icons if custom_icons else []
self._interval: float = interval
self._position: Literal["start", "end"] = position
self._indent: int = indent
self._spinner_color: Color = Color(spinner_color) if spinner_color else Color()
self._done_color: Color = Color(done_color) if done_color else Color()
# Terminal dimensions
self._term_columns: int = _get_terminal_width()
self._term_lines: int = _get_terminal_height()
# Threading
self._stop_e = Event() # Create an event that will be used to stop the spinner
self._is_running: bool = False
self._process = None
# Setup values
self.update(text=self._text,
icons_id=self._icons_id,
done_icon=self._done_icon,
custom_icons=self._icons,
interval=self._interval,
position=self._position,
indent=self._indent,
spinner_color=self._spinner_color,
done_color=self._done_color)
[docs]
def start(self, text: str = ''):
"""
Start the spinner
:param text: Text to display next to the spinner. Omit for no text.
"""
if self._is_running:
# Already running
return
self._is_running = True
if self._stop_e.is_set():
# Reset event
self._stop_e.clear()
if text:
# Update text
self._text = ShortenText(text=text).__str__()
# Hide cursor
stdout.write(ANSIEscape.HIDE_CURSOR)
stdout.flush()
# Start thread
self._process = Process(target=self._run)
self._process.start()
[docs]
def stop(self,
done_icon: str | None = None,
done_color: str | tuple[int, int, int] | Color | None = None):
"""
Stop the spinner
Optionally, allow custom done icon and color for the current spinner.
:param done_icon: Icon to use for when the spinner is complete
:param done_color: Color for the done icon
"""
if self._process is None:
# Not running
return
self._is_running = False
# Stop process
self._stop_e.set()
self._process.join()
# Cleanup
_done_color: Color = Color(done_color) if done_color else self._done_color
self._cleanup(done_icon=done_icon, done_color=_done_color.tc_color)
# Show cursor
stdout.write(ANSIEscape.SHOW_CURSOR)
stdout.flush()
# Ensure enough room
print()
[docs]
def update(self, text: str | AnyStatus | AnyLabeledData | None = None,
icons_id: str | None = None,
done_icon: str | None = None,
custom_icons: list[str] | None = None,
interval: float | None = None,
position: Literal["start", "end"] | None = None,
indent: int | None = None,
spinner_color: str | tuple[int, int, int] | Color | None = None,
done_color: str | tuple[int, int, int] | Color | None = None):
"""
Update values for the spinner.
Some values can only be updated when the spinner is not running. See documentation for more details.
**Position**
- ``"start"``: The spinner will be placed at the start of the line.
- ``"end"``: The spinner will be placed at the end of the line.
:param text: Text to display next to the spinner. Omit for no text.
:param icons_id: ID string of one of the built-in icon sets from ``Spinner.BUILTIN_ICONS``
:param done_icon: Icon to use for when the spinner is complete
:param custom_icons: A list of strings to use as the spinner icon. Overrides ``icons_id``
:param interval: The delay between spinner icon updates.
:param position: Where the spinner should be placed on a line.
:param indent: How many spaces the line should be indented.
:param spinner_color: Color of the spinner icon while active.
:param done_color: Color for the spinner icon when stopped.
"""
if not self._is_running:
# Values not allowed to be changed while running
self._icons_id = icons_id if icons_id else self._icons_id
self._interval = interval if interval else self._interval
self._position = position if position else self._position
self._indent = indent if indent else self._indent
# Update colors
if spinner_color:
self._spinner_color = Color(spinner_color)
if done_color:
self._done_color = Color(done_color)
# Setup icons
if custom_icons:
# Custom icons
self._icons = [icon[0] for icon in custom_icons]
elif icons_id:
# Built-in icons
if icons_id in list(self.BUILTIN_ICONS.keys()):
self._icons = self.BUILTIN_ICONS[icons_id]
else:
raise ValueError(f"Invalid Icon ID: {icons_id}\nValid Icon IDs: {list(self.BUILTIN_ICONS.keys())}")
else:
# Default icons
self._icons = self.BUILTIN_ICONS["default"]
# Values allowed to be changed while running
self._done_icon = done_icon if done_icon else self._done_icon
# Resolve text
self._text = ShortenText(text=text).__str__() if text else self._text
def _cleanup(self,
done_icon: str | None = None,
done_color: str | tuple[int, int, int] | Color | None = None):
"""
Remove or replace the spinner based on position
:param done_icon: Icon to use for when the spinner is complete
:param done_color: Color for the spinner icon when stopped.
"""
stdout.write(ANSIEscape.ERASE_LINE)
_done_color = Color(done_color) if done_color else self._done_color
_done_icon = done_icon if done_icon else self._done_icon
# Get icon, then redraw spinner
icon: str = ''
if self._position == "start":
if _done_icon:
icon = _format_col_text(text=f'{_done_icon[0]} ', color=_done_color.tc_color)
stdout.write(f"\r{' ' * self._indent}{icon}{self._text}")
elif self._position == "end":
if _done_icon:
icon = _format_col_text(text=f' {_done_icon[0]}', color=_done_color.tc_color)
stdout.write(f"\r{' ' * self._indent}{self._text}{icon}")
def _run(self):
"""
Function for ``self._process`` to run and display the spinner
"""
try:
itr: int = 0
while not self._stop_e.is_set():
# Get icon
cur_icon: str = _format_col_text(text=self._icons[itr], color=self._spinner_color.tc_color)
# Print spinner
if self._position == "start":
stdout.write(f"\r{' ' * self._indent}{cur_icon} {self._text}")
elif self._position == "end":
stdout.write(f"\r{' ' * self._indent}{self._text} {cur_icon}")
# Move to next icon
itr = (itr + 1) % len(self._icons)
sleep(self._interval)
except (EOFError, KeyboardInterrupt):
# Clean up and show cursor
self._cleanup(done_icon='✘', done_color="red")
stdout.write(ANSIEscape.SHOW_CURSOR)
stdout.flush()