graceful_interrupt_handler.py

import errno
import os
import signal
import unittest
import time


class graceful_interrupt_handler(object):
    '''
    Usage:
    with graceful_interrupt_handler as handler:
        do_stuff()
        if handler.interrupted:
            do_more_stuff()

    handler.interrupted is called immediately upon receiving the termination
    signal, so do_more_stuff() in the above example should be something that
    allows do_stuff to finish cleanly. The script will exit after
    handler.interrupted finishes unless you wrap calls in itself, at which
    point it will require as many consecutive calls to kill as you wrap.

    A fully tested suite based off http://stackoverflow.com/a/10972804/2638784
    '''

    def __init__(self, signals=(signal.SIGINT, signal.SIGTERM)):
        self.signals = signals
        self.original_handlers = {}

    def __enter__(self):
        self.interrupted = False
        self.released = False

        for sig in self.signals:
            self.original_handlers[sig] = signal.getsignal(sig)
            signal.signal(sig, self.handler)

        return self

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

    def __exit__(self, type, value, tb):
        self.release()

    def release(self):
        if self.released:
            return False

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

        self.released = True
        return True

#################################################
# TESTING
#################################################

TESTFN = "$-testfile"
TESTFN2 = "$-testfile2"
POSIX = os.name == 'posix'
WINDOWS = os.name == 'nt'
test_files = []


def safe_remove(file):
    "Convenience function for removing temporary test files"
    try:
        os.remove(file)
    except OSError as err:
        if err.errno != errno.ENOENT:
            raise


@unittest.skipIf(WINDOWS, "Unix-only test!")
class GracefulInterruptHandlerTestCase(unittest.TestCase):

    def test_exit_cleanly(self):
        # Make sure handler is called on SIGTERM

        pid = os.fork()
        if pid:
            # we're the parent process
            time.sleep(1)
            os.kill(pid, signal.SIGTERM)
            time.sleep(1)
            self.assertTrue(os.path.exists(TESTFN))
        else:
            # we're the child process
            with graceful_interrupt_handler() as h:
                time.sleep(5)
                if h.interrupted:
                    open(TESTFN, 'wb').write('yo')

    def test_exit_cleanly_with_SIGKILL(self):
        # Make sure handler is not called on SIGKILL

        pid = os.fork()
        if pid:
            # we're the parent process
            time.sleep(1)
            os.kill(pid, signal.SIGKILL)
            time.sleep(1)
            self.assertFalse(os.path.exists(TESTFN2))
        else:
            # we're the child process
            with graceful_interrupt_handler() as h:
                time.sleep(5)
                if h.interrupted:
                    # will never get created
                    open(TESTFN2, 'wb').write('yo')


if __name__ == '__main__':
    safe_remove(TESTFN)
    unittest.main(verbosity=2)
    safe_remove(TESTFN)