Context manager

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

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,13 +37,15 @@ 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)
 
 
 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()