DelayedInterrupt class - delaying the handling of process signals in Python

delayedinterrupt.py

import signal
import logging


# class based on: http://stackoverflow.com/a/21919644/487556
class DelayedInterrupt(object):
    def __init__(self, signals):
        if not isinstance(signals, list) and not isinstance(signals, tuple):
            signals = [signals]
        self.sigs = signals        

    def __enter__(self):
        self.signal_received = {}
        self.old_handlers = {}
        for sig in self.sigs:
            self.signal_received[sig] = False
            self.old_handlers[sig] = signal.getsignal(sig)
            def handler(s, frame):
                self.signal_received[sig] = (s, frame)
                # Note: in Python 3.5, you can use signal.Signals(sig).name
                logging.info('Signal %s received. Delaying KeyboardInterrupt.' % sig)
            self.old_handlers[sig] = signal.getsignal(sig)
            signal.signal(sig, handler)

    def __exit__(self, type, value, traceback):
        for sig in self.sigs:
            signal.signal(sig, self.old_handlers[sig])
            if self.signal_received[sig] and self.old_handlers[sig]:
                self.old_handlers[sig](*self.signal_received[sig])

test_delayedinterrupt.py

import os
import signal
from delayedinterrupt import DelayedInterrupt
from mock import Mock


def test_delayed_interrupt_with_one_signal():
    # check behavior without DelayedInterrupt
    a = Mock()
    b = Mock()
    c = Mock()
    try:
        a()
        os.kill(os.getpid(), signal.SIGINT)
        b()
    except KeyboardInterrupt:
        c()
    a.assert_called_with()
    b.assert_not_called()
    c.assert_called_with()

    # test behavior with DelayedInterrupt
    a = Mock()
    b = Mock()
    c = Mock()
    try:
        with DelayedInterrupt(signal.SIGINT):
            a()
            os.kill(os.getpid(), signal.SIGINT)
            b()
    except KeyboardInterrupt:
        c()
    a.assert_called_with()
    b.assert_called_with()
    c.assert_called_with()


def test_delayed_interrupt_with_multiple_signals():
    a = Mock()
    b = Mock()
    c = Mock()
    try:
        with DelayedInterrupt([signal.SIGTERM, signal.SIGINT]):
            a()
            os.kill(os.getpid(), signal.SIGINT)
            os.kill(os.getpid(), signal.SIGTERM)
            b()
    except KeyboardInterrupt:
        c()
    a.assert_called_with()
    b.assert_called_with()
    c.assert_called_with()