blob: 491e7e8d33fbccb2f055c948faf688aab15ab026 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
|
"""A simple spinner module"""
import itertools
import sys
import threading
import time
class Spinner:
"""A simple spinner class"""
def __init__(
self,
message: str = "Loading...",
delay: float = 0.1,
plain_output: bool = False,
) -> None:
"""Initialize the spinner class
Args:
message (str): The message to display.
delay (float): The delay between each spinner update.
plain_output (bool): Whether to display the spinner or not.
"""
self.plain_output = plain_output
self.spinner = itertools.cycle(["-", "/", "|", "\\"])
self.delay = delay
self.message = message
self.running = False
self.spinner_thread = None
def spin(self) -> None:
"""Spin the spinner"""
if self.plain_output:
self.print_message()
return
while self.running:
self.print_message()
time.sleep(self.delay)
def print_message(self):
sys.stdout.write(f"\r{' ' * (len(self.message) + 2)}\r")
sys.stdout.write(f"{next(self.spinner)} {self.message}\r")
sys.stdout.flush()
def __enter__(self):
"""Start the spinner"""
self.running = True
self.spinner_thread = threading.Thread(target=self.spin)
self.spinner_thread.start()
return self
def __exit__(self, exc_type, exc_value, exc_traceback) -> None:
"""Stop the spinner
Args:
exc_type (Exception): The exception type.
exc_value (Exception): The exception value.
exc_traceback (Exception): The exception traceback.
"""
self.running = False
if self.spinner_thread is not None:
self.spinner_thread.join()
sys.stdout.write(f"\r{' ' * (len(self.message) + 2)}\r")
sys.stdout.flush()
def update_message(self, new_message, delay=0.1):
"""Update the spinner message
Args:
new_message (str): New message to display.
delay (float): The delay in seconds between each spinner update.
"""
self.delay = delay
self.message = new_message
if self.plain_output:
self.print_message()
|
