signal_handling.py

import signal
import logging
import time
import sys

from contextlib import contextmanager
from typing import Callable


LOG = logging.getLogger(__name__)


class SignalInfo:
    caught = False

    def __init__(self, sig):
        self.signal = sig


@contextmanager
def catch(signalnum: int, handler: Callable = None) -> SignalInfo:
    """Catch signals with some sanity"""
    info = SignalInfo(signalnum)

    def _handler(*args):
        info.caught = True

        if handler:
            handler(*args)
        else:
            LOG.warning("Received signal {}, waiting until pending tasks are finished...".format(signalnum))

    original_handler = signal.getsignal(signalnum)
    signal.signal(signalnum, _handler)

    yield info

    signal.signal(signalnum, original_handler)


def quit(*_):
    LOG.info("Received SIGINT, Ciao!")
    sys.exit(0)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)

    while True:
        # Program would wait for sleep to finish on SIGINT
        with catch(signal.SIGINT) as sigint:
            LOG.info("Doing something that can't be interrupted.")
            time.sleep(5)

            if sigint.caught:
                quit()

        # Program will exit immediately on SIGINT
        with catch(signal.SIGINT, handler=quit):
            LOG.info("Doing something that can be interrupted safely.")
            time.sleep(5)