Progress Spinner Using Objects and Generators
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import sys | |
import time | |
spinners = { | |
"angles": ["◢", "◣", "◤", "◥"], | |
"circle": ["◜", "◠", "◝", "◞", "◡", "◟"], | |
"dots": ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"], | |
"ticks": ["-", "\\", "|", "/"], | |
} | |
class Spinner: | |
def __init__(self, style=None, text=None): | |
if style is None: | |
style = "ticks" | |
if text is None: | |
text = "Processing..." | |
self.style = style | |
self.text = text | |
self._error = False | |
self.spinner = self.frame_generator() | |
def __enter__(self): | |
return self | |
def __exit__(self, exc_type, exc_value, traceback): | |
if exc_type is not None: | |
sys.stdout.write("\r𝘅\n") | |
else: | |
sys.stdout.write("\r✔\n") | |
self.spinner.close() | |
self._cleanup() | |
def frame_generator(self): | |
frames = spinners[self.style] | |
i = 0 | |
while True: | |
yield frames[i % len(frames)] | |
i += 1 | |
def update(self, text=None): | |
if text is not None: | |
self.text = text | |
sys.stdout.write(f"\r{next(self.spinner)} {self.text}") | |
def error(self, error): | |
self._error = True | |
sys.stdout.write("\r𝘅\n") | |
raise | |
def close(self): | |
if not self._error: | |
sys.stdout.write("\r✔\n") | |
self.spinner.close() | |
self._cleanup() | |
def __repr__(self): | |
return f"{self.__class__.__name__}('{self.style}')" | |
def __str__(self): | |
return f"{next(self.spinner)} {self.text}" | |
@staticmethod | |
def _cleanup(): | |
sys.stdout.flush() | |
if __name__ == "__main__": | |
print("** TESTING AS CONTEXT MANAGER **") | |
with Spinner() as spinner: | |
steps = 10 | |
for i in range(steps): | |
spinner.update() | |
time.sleep(0.25) | |
print("\n** TESTING AS CONTEXT MANAGER WITH CUSTOM PROMPTS **") | |
jobs = ["angles spinner", "circle spinner"] | |
for job in jobs: | |
style = job.split(" ")[0] if "exception" not in job else None | |
with Spinner(style=style) as spinner: | |
steps = 10 | |
for i in range(steps): | |
spinner.update(f"processing {job} (step {i + 1} of {steps})...") | |
time.sleep(0.25) | |
print("\n** TESTING AS OBJECT WITH DEFINED TEXT **") | |
jobs = ["dots spinner", "spinner with exception"] | |
for job in jobs: | |
style = job.split(" ")[0] if "exception" not in job else None | |
s = Spinner(style=style, text=job) | |
steps = 10 | |
try: | |
for i in range(steps): | |
s.update() | |
time.sleep(0.25) | |
if "exception" in job and i == 6: | |
raise ValueError("test exception") | |
except Exception as e: | |
s.error(e) | |
finally: | |
s.close() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment