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