aboutsummaryrefslogtreecommitdiff
path: root/autogpts/autogpt/autogpt/app/spinner.py
diff options
context:
space:
mode:
Diffstat (limited to 'autogpts/autogpt/autogpt/app/spinner.py')
-rw-r--r--autogpts/autogpt/autogpt/app/spinner.py70
1 files changed, 70 insertions, 0 deletions
diff --git a/autogpts/autogpt/autogpt/app/spinner.py b/autogpts/autogpt/autogpt/app/spinner.py
new file mode 100644
index 000000000..8b2aa6c3c
--- /dev/null
+++ b/autogpts/autogpt/autogpt/app/spinner.py
@@ -0,0 +1,70 @@
+"""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 start(self):
+ self.running = True
+ self.spinner_thread = threading.Thread(target=self.spin)
+ self.spinner_thread.start()
+
+ def stop(self):
+ 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 __enter__(self):
+ """Start the spinner"""
+ self.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.stop()