ターミナル上で長めの処理をするときに使いたい、くっそ便利なぐるぐるクラス[Python]

デモ

こんな感じのぐるぐるを簡単に作れます。

コード

class Spinner:
    def __init__(
        self,
        text: str = "Please wait...",
        etext: str = "",
        overwrite: bool = True,
    ) -> None:
        self.text = text
        self.end_text = etext
        self.overwrite = overwrite

    def start(self) -> None:
        self._stop_flag = False
        self._spinner_thread = threading.Thread(target=self._spinner)
        self._spinner_thread.setDaemon(True)
        self._spinner_thread.start()

    def stop(self, etext: str = "", overwrite: bool = True) -> None:
        if self._spinner_thread and self._spinner_thread.is_alive():
            self._stop_flag = True
            self._spinner_thread.join()

        text = etext or self.end_text
        overwrite = self.overwrite if not self.overwrite else overwrite

        if overwrite:
            if etext == "":
                print(f"\r\033[2K\033[G", end="")
            else:
                print(f"\r\033[2K\033[G{etext}")
        else:
            if etext == "":
                print(f"\033[1D\033[K\n", end="")
            else:
                print(f"\033[1D\033[K\n{etext}")

    def _spinner(self) -> None:
        chars = itertools.cycle(r"/-\|")
        while not self._stop_flag:
            print(f"\r{self.text} {next(chars)}", end="")
            time.sleep(0.2)

    def __enter__(self) -> None:
        self.start()

    def __exit__(self, exception_type, exception_value, traceback) -> None:
        self.stop()

    def __call__(self, func: Callable) -> Callable:
        @functools.wraps(func)
        def wrapped(*args, **kwargs):
            with self:
                return func(*args, **kwargs)

        return wrapped

使い方

textには、ぐるぐるを付与したい文、etextには処理終了時に表示したい文を指定します。

3つの使い方があります。

start()とstop()の間に処理を書く方法

spinner = Spinner(text="Loading...", etxet="Loading... Done.")
spinner.start()

time.sleep(3)

spinner.stop()

withで囲む方法

with Spinner("Loading...", "Loading... Done."):
    time.sleep(3)

デコレータとして使う方法

@Spinner("Loading...", "Loading... Done.")
def my_func():
    time.sleep(3)

my_func()

実行結果はどれも同じです。

overwrite=Falseで処理後の出力を次の行に指定することができます。

spinner = Spinner("Loading...", "Done.", overwrite=False)
spinner.start()

time.sleep(3)

spinner.stop()

実用例

start()とstop()の間に処理を書く方法では、stop()にetextを指定できるので、中の処理結果に応じた出力ができます。

spinner = Spinner("当たるかな...")
spinner.start()

time.sleep(3)
n = np.random.randint(0, 2)
result = {0: "当たり！", 1: "ハズレ..."}

spinner.stop(etext=result[n], overwrite=False)

運が悪かったようです。

終わりに

最近cliの開発をしているときにあったら便利だな〜ということで作成しました。
みなさまも是非お使いください。

参考にした記事等

2022/7/22
一部コード修正