# :copyright: (c) 2021 by Pavlo Dmytrenko. # :license: MIT, see LICENSE for more details. """ yaspin.core ~~~~~~~~~~~ A lightweight terminal spinner. """ from __future__ import annotations from collections.abc import Callable, Generator, Iterator, Sequence from contextlib import contextmanager from dataclasses import dataclass from datetime import timedelta from typing import ( Any, cast, Final, Protocol, runtime_checkable, TextIO, TYPE_CHECKING, TypeVar, ) import functools import itertools import shutil import signal import sys import threading import time import warnings from termcolor import ATTRIBUTES, colored, COLORS, HIGHLIGHTS from .constants import SPINNER_ATTRS if TYPE_CHECKING: from types import FrameType, TracebackType SignalHandlers = Callable[[int, FrameType | None], Any] | int | None Fn = TypeVar("Fn", bound=Callable[..., Any]) ENCODING: Final[str] = "utf-8" class SafeStreamWrapper: """A wrapper that handles closed streams gracefully.""" def __init__(self, stream: TextIO, warn_on_closed: bool = False) -> None: self._stream = stream self._warn_on_closed = warn_on_closed self._warned_already = False # Avoid warning spam def write(self, text: str) -> None: """Write to stream, optionally warning if stream is closed.""" if not self._stream.closed: self._stream.write(text) elif self._warn_on_closed and not self._warned_already: warnings.warn( "Attempted to write to closed stream. Output ignored. " "This may indicate a stream lifecycle management issue.", UserWarning, stacklevel=3, ) self._warned_already = True def flush(self) -> None: """Flush stream, silently ignoring if stream is closed.""" if not self._stream.closed: self._stream.flush() # Note: don't warn on flush - it is often called during cleanup def isatty(self) -> bool: """Check if stream is a TTY, returning False if closed.""" return not self._stream.closed and self._stream.isatty() @property def closed(self) -> bool: """Check if the underlying stream is closed.""" return self._stream.closed def __getattr__(self, name: str) -> Any: """Delegate other attributes to the underlying stream.""" return getattr(self._stream, name) def to_unicode(text_type: str | bytes, encoding: str = ENCODING) -> str: if isinstance(text_type, bytes): return text_type.decode(encoding) return text_type @dataclass class Spinner: frames: Sequence[str] interval: int default_spinner = Spinner("⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏", 80) @runtime_checkable class SignalHandlerProtocol(Protocol): def __call__(self, signum: int, frame: Any, spinner: Yaspin) -> None: ... def default_handler(signum: int, frame: Any, spinner: Yaspin) -> None: """Signal handler, used to gracefully shut down the ``spinner`` instance when specified signal is received by the process running the ``spinner``. ``signum`` and ``frame`` are mandatory arguments. Check ``signal.signal`` function for more details. """ spinner.fail() spinner.stop() sys.exit(0) def fancy_handler(signum: int, frame: Any, spinner: Yaspin) -> None: """Signal handler, used to gracefully shut down the ``spinner`` instance when specified signal is received by the process running the ``spinner``. ``signum`` and ``frame`` are mandatory arguments. Check ``signal.signal`` function for more details. """ spinner.red.fail("✘") spinner.stop() sys.exit(0) class Yaspin: """Implements a context manager that spawns a thread to write spinner frames into a tty (stdout by default) during context execution. """ # When Python finds its output attached to a terminal, # it sets the sys.stdout.encoding attribute to the terminal's encoding. # The print statement's handler will automatically encode unicode # arguments into bytes. def __init__( self, spinner: Spinner = default_spinner, text: str = "", color: str | None = None, on_color: str | None = None, attrs: Sequence[str] | None = None, reversal: bool = False, side: str = "left", sigmap: dict[signal.Signals, SignalHandlers] | None = None, timer: bool = False, ellipsis: str = "", stream: TextIO | None = None, warn_on_closed_stream: bool = False, ) -> None: # Stream raw_stream = stream or sys.stdout self._stream = SafeStreamWrapper(raw_stream, warn_on_closed=warn_on_closed_stream) self._stream_lock = threading.Lock() # Spinner self._spinner = self._set_spinner(spinner) self._frames = self._set_frames(self._spinner, reversal) self._interval = self._set_interval(self._spinner) self._cycle = self._set_cycle(self._frames) # Color Specification self._color = self._set_color(color) if color else color self._on_color = self._set_on_color(on_color) if on_color else on_color self._attrs = self._set_attrs(attrs) if attrs else set() self._color_func = self._compose_color_func() # Other self._text = text self._side = self._set_side(side) self._reversal = reversal self._timer = timer self._ellipsis = ellipsis self._terminal_width: int = shutil.get_terminal_size().columns self._start_time: float | None = None self._stop_time: float | None = None # Helper flags self._stop_spin: threading.Event | None = None self._hide_spin: threading.Event | None = None self._spin_thread: threading.Thread | None = None self._last_frame: str | None = None self._hidden_level = 0 self._cur_line_len = 0 # Signals self._sigmap = sigmap if sigmap else {} # Maps signals to their default handlers in order to reset # custom handlers set by ``sigmap`` at the cleanup phase. self._dfl_sigmap: dict[signal.Signals, SignalHandlers] = {} # Dunders # def __repr__(self) -> str: return f"" def __enter__(self) -> Yaspin: self.start() return self def __exit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None, ) -> None: if self._spin_thread is None: raise RuntimeError("spin thread is None") # Avoid stop() execution for the 2nd time if self._spin_thread.is_alive(): self.stop() def __call__(self, fn: Fn) -> Fn: @functools.wraps(fn) def inner(*args: Any, **kwargs: Any) -> Fn: with self: return fn(*args, **kwargs) return cast(Fn, inner) def __getattr__(self, name: str) -> Yaspin: # CLI spinners if name in SPINNER_ATTRS: from .spinners import Spinners sp = getattr(Spinners, name) self.spinner = sp # Color Attributes: "color", "on_color", "attrs" elif name in set(key for d in [ATTRIBUTES, COLORS, HIGHLIGHTS] for key in d): # Call appropriate property setters; # _color_func is updated automatically by setters. if name in ATTRIBUTES: self.attrs = [name] # calls property setter if name in COLORS: self.color = name # calls property setter if name in HIGHLIGHTS: self.on_color = name # calls property setter # Side: "left" or "right" elif name in ("left", "right"): self.side = name # calls property setter # Common error for unsupported attributes else: raise AttributeError(f"'{self.__class__.__name__}' object has no attribute: '{name}'") return self # Properties # @property def spinner(self) -> Spinner: return self._spinner @spinner.setter def spinner(self, sp: Spinner) -> None: self._spinner = self._set_spinner(sp) self._frames = self._set_frames(self._spinner, self._reversal) self._interval = self._set_interval(self._spinner) self._cycle = self._set_cycle(self._frames) @property def text(self) -> str: return self._text @text.setter def text(self, txt: str) -> None: self._text = txt @property def color(self) -> str | None: return self._color @color.setter def color(self, value: str) -> None: self._color = self._set_color(value) if value else value self._color_func = self._compose_color_func() # update @property def on_color(self) -> str | None: return self._on_color @on_color.setter def on_color(self, value: str) -> None: self._on_color = self._set_on_color(value) if value else value self._color_func = self._compose_color_func() # update @property def attrs(self) -> Sequence[str]: return list(self._attrs) @attrs.setter def attrs(self, value: Sequence[str]) -> None: new_attrs = self._set_attrs(value) if value else set() self._attrs = self._attrs.union(new_attrs) self._color_func = self._compose_color_func() # update @property def side(self) -> str: return self._side @side.setter def side(self, value: str) -> None: self._side = self._set_side(value) @property def ellipsis(self) -> str: return self._ellipsis @ellipsis.setter def ellipsis(self, value: str) -> None: self._ellipsis = value @property def reversal(self) -> bool: return self._reversal @reversal.setter def reversal(self, value: bool) -> None: self._reversal = value self._frames = self._set_frames(self._spinner, self._reversal) self._cycle = self._set_cycle(self._frames) @property def elapsed_time(self) -> float: if self._start_time is None: return 0 if self._stop_time is None: return time.time() - self._start_time return self._stop_time - self._start_time # Public # def start(self) -> None: """ Start the spinner animation in a separate thread. Initializes and starts the spinner animation by hiding the cursor, recording the start time, and creating a new thread to run the spinner. It also sets up necessary threading events to control the spinner's behavior. If signal handlers are registered, they will be set up before starting the spinner. In case of any failure that prevents the spinner from starting, the cursor will be shown to ensure it is not left hidden. """ if self._sigmap: self._register_signal_handlers() self._hide_cursor() self._start_time = time.time() # Reset value to properly calculate subsequent spinner starts (if any) self._stop_time = None self._stop_spin = threading.Event() self._hide_spin = threading.Event() self._spin_thread = threading.Thread(target=self._spin) try: self._spin_thread.start() finally: # Ensure cursor is not hidden if any failure occurs that prevents # getting it back self._show_cursor() def stop(self) -> None: """ Stops the spinner and performs necessary cleanup. Records the stop time, resets signal handlers to their default if they were modified, stops the spinning thread, clears the spinner line, and shows the cursor. Raises: RuntimeError: If the stop_spin event is None. """ self._stop_time = time.time() if self._dfl_sigmap: # Reset registered signal handlers to default ones self._reset_signal_handlers() if self._spin_thread is not None: if self._stop_spin is None: raise RuntimeError("stop_spin event is None") self._stop_spin.set() self._spin_thread.join() self._clear_line() self._show_cursor() def hide(self) -> None: """ Hide the spinner to allow for custom writing to the terminal. Sets a flag to indicate that the spinner should be hidden, clears the current line in the terminal, and flushes the stdout buffer. It ensures that the spinner thread is alive and the hide flag is not already set before performing these actions. Raises: RuntimeError: If the hide_spin attribute is None. """ thr_is_alive = self._spin_thread and self._spin_thread.is_alive() if self._hide_spin is None: raise RuntimeError("hide_spin is None") if thr_is_alive and not self._hide_spin.is_set(): with self._stream_lock: # set the hidden spinner flag self._hide_spin.set() self._clear_line() # flush the stream buffer so the current line # can be rewritten to self._stream.flush() @contextmanager def hidden(self) -> Generator[None, None, None]: """ Temporarily hides the spinner within a context block. This method can be nested. When the context is entered, the spinner is hidden if it is not already hidden. When the context is exited, the spinner is shown again if it was hidden by this method. Yields: None: This method is a generator that yields control back to the caller. """ if self._hidden_level == 0: self.hide() self._hidden_level += 1 try: yield finally: self._hidden_level -= 1 if self._hidden_level == 0: self.show() def show(self) -> None: """ Show the hidden spinner. Checks if the spinner thread is alive and if the spinner is currently hidden. If both conditions are met, it clears the hidden spinner flag and clears the current line to ensure the spinner is not appended to it. Raises: RuntimeError: If the `_hide_spin` attribute is `None`. """ thr_is_alive = self._spin_thread and self._spin_thread.is_alive() if self._hide_spin is None: raise RuntimeError("hide_spin is None") if thr_is_alive and self._hide_spin.is_set(): with self._stream_lock: # clear the hidden spinner flag self._hide_spin.clear() # clear the current line so the spinner is not appended to it self._clear_line() def write(self, text: str) -> None: """ Write text in the terminal without breaking the spinner. Ensures that the spinner is temporarily cleared, the text is written to the terminal, and then the spinner is restored. Args: text (str): The text to be written to the terminal. """ # similar to tqdm.write() # https://pypi.python.org/pypi/tqdm#writing-messages with self._stream_lock: self._clear_line() _text = to_unicode(text) if isinstance(text, str | bytes) else str(text) self._stream.write(f"{_text}\n") self._cur_line_len = 0 def ok(self, text: str = "OK") -> None: """Set Ok (success) finalizer to a spinner.""" _text = text if text else "OK" self._freeze(_text) def fail(self, text: str = "FAIL") -> None: """Set fail finalizer to a spinner.""" _text = text if text else "FAIL" self._freeze(_text) def _supports_ansi_codes(self) -> bool: """Check if the stream supports ANSI escape sequences. Returns True for TTY streams where ANSI codes work properly, False for non-TTY streams (files, StringIO, etc.) where ANSI codes would appear as literal text. """ return self._stream.isatty() def is_jupyter(self) -> bool: warnings.warn( "is_jupyter() is deprecated and misleading. " "It detects non-TTY streams, not Jupyter environments. " "Use _supports_ansi_codes() instead.", DeprecationWarning, stacklevel=2, ) return not self._supports_ansi_codes() # Protected # def _freeze(self, final_text: str) -> None: """ Stop the spinner and display the final frame. Stops the spinner, composes the last frame with the provided final text, and 'freezes' it by writing the final frame to the standard output. Args: final_text (str): The final text to be displayed when the spinner stops. Raises: RuntimeError: If the last frame is None. """ text = to_unicode(final_text) self._last_frame = self._compose_out(text, mode="last") # Should be stopped here, otherwise prints after # self._freeze call will mess up the spinner self.stop() with self._stream_lock: if self._last_frame is None: raise RuntimeError("last_frame is None") self._stream.write(self._last_frame) self._cur_line_len = 0 def _spin(self) -> None: """ Handles the spinning animation. Continuously updates the spinner's output on the terminal until the `_stop_spin` event is set. If the `_hide_spin` event is set, it temporarily pauses the spinning. Raises: RuntimeError: If `_stop_spin` is None. """ if self._stop_spin is None: raise RuntimeError("stop_spin is None") while not self._stop_spin.is_set(): if self._hide_spin is not None and self._hide_spin.is_set(): # Wait a bit to avoid wasting cycles time.sleep(self._interval) continue # Compose output spin_phase = next(self._cycle) out = self._compose_out(spin_phase) # Write with self._stream_lock: self._clear_line() self._stream.write(out) self._stream.flush() self._cur_line_len = max(self._cur_line_len, len(out)) # Wait self._stop_spin.wait(self._interval) def _compose_color_func(self) -> Callable[..., str] | None: """ Compose a color function based on the current environment. If the environment is Jupyter, returns None as ANSI color control sequences are problematic in Jupyter notebooks. Otherwise, returns a partial function that applies the specified color, background color, and attributes to text. """ if not self._supports_ansi_codes(): # ANSI Color Control Sequences are problematic in non-TTY streams return None return functools.partial( colored, color=self._color, on_color=self._on_color, attrs=list(self._attrs), ) def _compose_out(self, frame: str, mode: str | None = None) -> str: """ Compose the output string for the spinner. Args: frame (str): The current frame of the spinner animation. mode (str, optional): The mode in which the output is generated. If None, the output is generated on the same line with a carriage return. If a value is provided, the output is generated on a new line. Returns: str: The composed output string including the spinner frame, text, timer, and any specified colors and positions. Raises: ValueError: If the terminal size is too small to display the spinner with the given settings. """ text = str(self._text) # Timer if self._timer: sec, fsec = divmod(round(100 * self.elapsed_time), 100) timer = f" ({timedelta(seconds=sec)}.{fsec:02.0f})" else: timer = "" # Truncate max_text_len = self._get_max_text_length(len(frame), len(timer)) if max_text_len < 1: raise ValueError( f"Terminal size {self._terminal_width} is too small to display spinner " "with the given settings." ) text = text[:max_text_len] + self._ellipsis if len(text) > max_text_len else text # Colors if self._color_func is not None: frame = self._color_func(frame) # Position if self._side == "right": frame, text = text, frame # Mode out = f"\r{frame} {text}{timer}" if mode is None else f"{frame} {text}{timer}\n" return out def _get_max_text_length(self, frame_width: int, timer_width: int) -> int: """ Calculate the maximum length of text that can be displayed within the terminal width. Args: frame_width (int): The width of the frame. timer_width (int): The width of the timer. Returns: int: The maximum length of text that can be displayed. """ ellipsis_width = len(self._ellipsis) # There is always a space between frame and text frame_width += 1 return self._terminal_width - frame_width - timer_width - ellipsis_width def _register_signal_handlers(self) -> None: """ Registers custom signal handlers for the spinner. Sets up signal handlers defined in the `_sigmap` attribute. It ensures that SIGKILL is not included. For each signal in `_sigmap`, stores the default signal handler for later restoration during the cleanup phase. Raises: ValueError: If an attempt is made to set a handler for the SIGKILL signal. """ # SIGKILL cannot be caught or ignored, and the receiving # process cannot perform any clean-up upon receiving this # signal. if signal.SIGKILL in self._sigmap: raise ValueError( "Trying to set handler for SIGKILL signal. " "SIGKILL cannot be caught or ignored in POSIX systems." ) for sig, sig_handler in self._sigmap.items(): # A handler for a particular signal, once set, remains # installed until it is explicitly reset. Store default # signal handlers for subsequent reset at cleanup phase. dfl_handler = signal.getsignal(sig) self._dfl_sigmap[sig] = dfl_handler # ``signal.SIG_DFL`` and ``signal.SIG_IGN`` are also valid # signal handlers and are not callables. if callable(sig_handler) and isinstance(sig_handler, SignalHandlerProtocol): # ``signal.signal`` accepts handler function which is # called with two arguments: signal number and the # interrupted stack frame. ``functools.partial`` solves # the problem of passing spinner instance into the handler # function. sig_handler = functools.partial(sig_handler, spinner=self) signal.signal(sig, sig_handler) def _reset_signal_handlers(self) -> None: """Resets the signal handlers to their default values.""" for sig, sig_handler in self._dfl_sigmap.items(): signal.signal(sig, sig_handler) def _hide_cursor(self) -> None: if self._stream.isatty(): # ANSI Control Sequence DECTCEM 1 does not work in Jupyter self._stream.write("\033[?25l") self._stream.flush() def _show_cursor(self) -> None: if self._stream.isatty(): # ANSI Control Sequence DECTCEM 2 does not work in Jupyter self._stream.write("\033[?25h") self._stream.flush() def _clear_line(self) -> None: if self._stream.isatty(): # ANSI Control Sequence EL does not work in Jupyter self._stream.write("\r\033[K") else: fill = " " * self._cur_line_len self._stream.write(f"\r{fill}\r") def _set_color(self, value: str) -> str: if not self._supports_ansi_codes(): Yaspin._warn_color_disabled() if value not in COLORS: raise ValueError( "'{}': unsupported color value. Use one of the: {}".format(value, ", ".join(COLORS.keys())) ) return value def _set_on_color(self, value: str) -> str: if not self._supports_ansi_codes(): Yaspin._warn_color_disabled() if value not in HIGHLIGHTS: raise ValueError( "'{}': unsupported on_color value. Use one of the: {}".format( value, ", ".join(HIGHLIGHTS.keys()) ) ) return value def _set_attrs(self, attrs: Sequence[str]) -> set[str]: if not self._supports_ansi_codes(): Yaspin._warn_color_disabled() for attr in attrs: if attr not in ATTRIBUTES: raise ValueError( "'{}': unsupported attribute value. Use one of the: {}".format( attr, ", ".join(ATTRIBUTES.keys()) ) ) return set(attrs) # Static # @staticmethod def _warn_color_disabled() -> None: warnings.warn( "color, on_color and attrs are not supported when output stream is not a TTY", stacklevel=3, ) @staticmethod def _set_spinner(spinner: Spinner) -> Spinner: if hasattr(spinner, "frames") and hasattr(spinner, "interval"): sp = default_spinner if not spinner.frames or not spinner.interval else spinner else: sp = default_spinner return sp @staticmethod def _set_side(side: str) -> str: if side not in ("left", "right"): raise ValueError("'{0}': unsupported side value. Use either 'left' or 'right'.") return side @staticmethod def _set_frames(spinner: Spinner, reversal: bool) -> str | Sequence[str]: """ Set the frames for the spinner, optionally reversing them. Args: spinner (Spinner): The spinner object containing the frames. reversal (bool): If True, the frames will be reversed. Returns: Union[str, Sequence[str]]: The frames to be used for the spinner. This can be a single string of frames or a sequence of frame strings. Raises: ValueError: If no frames are found in the spinner. """ uframes = None # unicode frames uframes_seq = None # sequence of unicode frames if isinstance(spinner.frames, str): uframes = spinner.frames # TODO (pavdmyt): support any type that implements iterable if isinstance(spinner.frames, Sequence): # Empty ``spinner.frames`` is handled by ``Yaspin._set_spinner`` if spinner.frames and isinstance(spinner.frames[0], bytes): uframes_seq = [to_unicode(frame) for frame in spinner.frames] else: uframes_seq = spinner.frames _frames = uframes or uframes_seq if not _frames: # Empty ``spinner.frames`` is handled by ``Yaspin._set_spinner``. # This code is very unlikely to be executed. However, it's still # here to be on a safe side. raise ValueError(f"{spinner!r}: no frames found in spinner") # Builtin ``reversed`` returns reverse iterator, # which adds unnecessary difficulty for returning # unicode value; # Hence using [::-1] syntax frames = _frames[::-1] if reversal else _frames return frames @staticmethod def _set_interval(spinner: Spinner) -> float: # Milliseconds to Seconds return spinner.interval * 0.001 @staticmethod def _set_cycle(frames: str | Sequence[str]) -> Iterator[str]: return itertools.cycle(frames)