InterruptHandler.py

"""
Contém utilitários para capturar e/ou indicar interrupções em processos.
"""

from __future__ import annotations

import signal
from typing import Any, Callable, Mapping, Sequence, Set


class InterruptHandler:
    """
    Classe utilitária que permite observar interrupções de processo, além de oferecer contextual.

    Os métodos desta classe NÃO são thread-safe. Certifique-se de instalar o handler uma única vez,
    e sempre a partir do main thread.
    """
    installed: bool = False
    interrupted = None
    signal_number = None
    signal_frame = None
    released = False
    signals: Set[int] = set(signal.SIGINT, signal.SIGTERM)

    _assign_signal_data = False
    _error_builders: Mapping[int, Callable[[int, Any], Exception]] = {
        [signal.SIGINT]: KeyboardInterrupt()
    }

    _main_handler = None

    def __init__(
        self,
        signals: Sequence[int] = None,
        errors: Mapping[int, Callable[[int, Any], Exception]] = None,
        auto_install = False,
        auto_release = False,
        assign_signal_data = False,
    ):
        self._auto_release = auto_release
        self._previous_handlers = {}
        self._assign_signal_data = assign_signal_data

        if signals and len(signals) > 0:
            self.signals = set(signals)

        if errors:
            self._error_builders = errors

        if auto_install:
            self.install()



    @staticmethod
    def enter(
        *signals: int,
        errors: Mapping[int, Callable[..., Exception]] = None,
    ):
        """Cria um `InterruptHandler` para uso contextual (i.e. com `with`).

        Exemplo::

            with InterruptHandler.enter() as handler:
                finished = False
                while not finished and not handler.interrupted:
                    # ...
                    # Realizar algum trabalho contínuo
                    # ...
                    finished = True

        Returns:
            Um `InterruptHandler`.
        """
        return InterruptHandler(
            signals = signals,
            errors = errors,
            auto_install = True,
            auto_release = True,
            assign_signal_data = True,
        )


    @property
    def main(self) -> InterruptHandler:
        """Retorna o handler global.

        Por padrão, não há um handler global, então `InterruptHandler.main is None`.
        Vide `InterruptHandler.global_install()` para inicializar um handler global.

        Returns:
            InterruptHandler: O handler global, ou `None` se nenhum tiver sido preparado.
        """
        return self._main_handler


    @classmethod
    def global_install(
        cls,
        *signals: int,
        errors: Mapping[int, Callable[[int, Any], Exception]] = None,
    ):
        """Configura um handler global.

        Este método deve ser chamado na main thread, uma única vez, antes de iniciar qualquer outro
        thread. Uma vez chamado, o método irá instalar um handler global. Todos os threads poderão
        então acessar a instância global de `InterruptHandler` via

        Args:
            signals (Sequence[int], optional): Os sinais a capturar/observar. Opcional; por padrão,
            observa SIGINT e SIGTERM.

            errors (Mapping[int, Callable[..., Exception]], optional): Factories de erros a serem
            lançados ao ser capturado um sinal.

        Raises:
            Exception: Caso o handler global já tenha sido lançado previamente.
        """
        if cls._main_handler:
            raise Exception('Global interrupt handler was already set up.')

        cls._main_handler = InterruptHandler(
            signals = signals,
            errors = errors,
            auto_install = True,
            auto_release = False,
            assign_signal_data = False,
        )


    def install(self) -> bool:
        """
        Configura signal handlers para serem tratados por este objeto.

        Mesmo se chamado mais de uma vez, a configuração dos handlers via `signal.signal()`
        acontecerá apenas uma vez, durante todo o tempo de vida de `InterruptHandler`.
        """
        if self.installed:
            return False

        self.interrupted = False
        self.released = False

        for sig in self.signals:
            self._previous_handlers[sig] = signal.getsignal(sig)
            signal.signal(sig, self._handler)

        return True


    def release(self):
        """
        Restaura os signal handlers observados para seus valores (handlers) originais.

        Este método não fará nada se:
        - `self.install()` não tiver sdo chamado previamente, direta ou indiretamente; e
        - `self.release()` já tiver sido chamado previamente, direta ou indiretamente.
        """
        if not self.installed or self.released:
            return False

        for sig in self.signals:
            signal.signal(sig, self._previous_handlers[sig])

        self.released = True
        return True


    def _handler(self, signum, frame):
        self.interrupted = True

        if self._auto_release:
            self.release()

        if self._assign_signal_data:
            self.signal_number = signum
            self.signal_frame = frame

        if signum in self._error_builders:
            raise self._error_builders[signum](signum, frame)


    def __enter__(self):
        self.install()


    def __exit__(self, *_args):
        self.release()