Progress Spinner Using Objects and Generators

spinner_object.py

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()