1# -*- coding: utf-8 -*-
2
3import sys
4import threading
5import time
6import itertools
7
8
9class Spinner(object):
10    spinner_cycle = itertools.cycle(u'⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏')
11
12    def __init__(self, beep=False, force=False):
13        self.beep = beep
14        self.force = force
15        self.stop_running = None
16        self.spin_thread = None
17
18    def start(self):
19        if sys.stdout.isatty() or self.force:
20            self.stop_running = threading.Event()
21            self.spin_thread = threading.Thread(target=self.init_spin)
22            self.spin_thread.start()
23
24    def stop(self):
25        if self.spin_thread:
26            self.stop_running.set()
27            self.spin_thread.join()
28
29    def init_spin(self):
30        while not self.stop_running.is_set():
31            next_val = next(self.spinner_cycle)
32            if sys.version_info[0] == 2:
33                next_val = next_val.encode('utf-8')
34            sys.stdout.write(next_val)
35            sys.stdout.flush()
36            time.sleep(0.07)
37            sys.stdout.write('\b')
38
39    def __enter__(self):
40        self.start()
41        return self
42
43    def __exit__(self, exc_type, exc_val, exc_tb):
44        self.stop()
45        if self.beep:
46            sys.stdout.write('\7')
47            sys.stdout.flush()
48        return False
49
50
51def spinner(beep=False, force=False):
52    """This function creates a context manager that is used to display a
53    spinner on stdout as long as the context has not exited.
54
55    The spinner is created only if stdout is not redirected, or if the spinner
56    is forced using the `force` parameter.
57
58    Parameters
59    ----------
60    beep : bool
61        Beep when spinner finishes.
62    force : bool
63        Force creation of spinner even when stdout is redirected.
64
65    Example
66    -------
67
68        with spinner():
69            do_something()
70            do_something_else()
71
72    """
73    return Spinner(beep, force)
74