BasaltMeter/BasaltTrainer/while_stopper.py

import sys

IS_WINDOWS: bool = sys.platform == "win32"

if IS_WINDOWS:
    import win32api
else:
    import signal


class WhileStopper:
    _is_registered: bool
    _stop_requested: bool

    def __init__(self) -> None:
        self._is_registered = False
        self._stop_requested = False

    def is_stop_requested(self) -> bool:
        if not self._is_registered:
            raise RuntimeError("unexpected stopper status")

        return self._stop_requested

    def register(self):
        if self._is_registered:
            raise RuntimeError("unexpected stopper status")

        self._stop_requested = False
        if IS_WINDOWS:
            win32api.SetConsoleCtrlHandler(_win_handler, True)
        else:
            signal.signal(signal.SIGINT, _posix_handler)
        self._is_registered = True

    def unregister(self):
        if not self._is_registered:
            raise RuntimeError("unexpected stopper status")

        if IS_WINDOWS:
            win32api.SetConsoleCtrlHandler(_win_handler, False)
        else:
            signal.signal(signal.SIGINT, signal.SIG_DFL)
        self._is_registered = False


INSTANCE: WhileStopper = WhileStopper()
"""The singleton of WhileStopper"""


def _win_handler(ctrl_type: int) -> bool:
    INSTANCE._stop_requested = True
    return True


def _posix_handler(signal, frame) -> None:
    INSTANCE._stop_requested = True