diff --git a/pyspin/spin.py b/pyspin/spin.py index 7093a0c..527c842 100644 --- a/pyspin/spin.py +++ b/pyspin/spin.py @@ -4,6 +4,7 @@ from __future__ import absolute_import, print_function import sys +import threading import time from functools import wraps from concurrent.futures import ThreadPoolExecutor @@ -14,7 +15,6 @@ else: text_type = str - Box1 = u'⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏' Box2 = u'⠋⠙⠚⠞⠖⠦⠴⠲⠳⠓' Box3 = u'⠄⠆⠇⠋⠙⠸⠰⠠⠰⠸⠙⠋⠇⠆' @@ -35,11 +35,14 @@ class Spinner(object): - - def __init__(self, frames): + def __init__(self, frames=Default, words="", ending="\n"): self.frames = frames self.length = len(frames) self.position = 0 + self.words = words + self.ending = ending + self.stop_running = None + self.spin_thread = None def current(self): return self.frames[self.position] @@ -52,6 +55,35 @@ def next(self): def reset(self): self.position = 0 + def start(self): + if sys.stdout.isatty(): + self.stop_running = threading.Event() + self.spin_thread = threading.Thread(target=self.init_spin) + self.spin_thread.start() + + def stop(self): + if self.spin_thread: + self.stop_running.set() + self.spin_thread.join() + + def init_spin(self): + while not self.stop_running.is_set(): + print(text_type("\r{0} {1}").format(self.next(), + self.words), + end="") + sys.stdout.flush() + time.sleep(0.1) + + def __enter__(self): + self.start() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() + print(self.ending, end="") + sys.stdout.flush() + return False + def make_spin(spin_style=Default, words="", ending="\n"): spinner = Spinner(spin_style) @@ -71,5 +103,7 @@ def wrapper(*args, **kwargs): print(ending, end="") return future.result() + return wrapper + return decorator diff --git a/test_pyspin.py b/test_pyspin.py index d0510a1..d981cfd 100644 --- a/test_pyspin.py +++ b/test_pyspin.py @@ -28,6 +28,7 @@ def test_make_spin(): @spin.make_spin(spin.Default, 'Downloading...') def fake_download(): time.sleep(2) + fake_download() @@ -36,6 +37,7 @@ def test_make_spin_with_args(): def fake_download(url, retry_times=3): print("Downloading {0}, will retry {1} times".format(url, retry_times)) time.sleep(2) + fake_download("https://www.example.com/text.txt", retry_times=5) @@ -43,6 +45,7 @@ def test_stop_on_exception(): @spin.make_spin(spin.Default, 'Downloading...') def fake_download(): 1 / 0 + try: fake_download() except ZeroDivisionError: @@ -53,7 +56,16 @@ def test_several_calls(): @spin.make_spin(spin.Default, 'Downloading...') def fake_download(): time.sleep(2) + print("Begin the first download.") fake_download() print("Begin the second download.") fake_download() + + +def test_context_manager(): + def fake_download(): + time.sleep(2) + + with spin.Spinner(): + fake_download()