A simple script to catch any valid OS signals and simulate an action in response to a signal

signals.py

import argparse
import functools
import logging
import os
import signal
import sys
import time

_LOG_FORMAT = "%(levelname)s:%(asctime)s:%(name)s:%(message)s"
logging.basicConfig(stream=sys.stdout, format=_LOG_FORMAT)
_logger = logging.getLogger()
_logger.setLevel(logging.INFO)


def cleanup(cleanup_time):
    _logger.info(f'Simulating cleanup for {cleanup_time} seconds.')
    time.sleep(cleanup_time)
    _logger.info('Cleanup complete.')


def signal_handler(cleanup_time, _signal, stack_frame):
    _logger.warning(f'Caught signal {signal.strsignal(_signal)}')

    # Simulate cleanup if the app is gracefully terminated
    if _signal == signal.SIGTERM:
        cleanup(cleanup_time)


def main(run_time):
    _logger.info(f'Sleeping for {run_time} seconds.')
    time.sleep(run_time)
    _logger.info('Sleep complete.')


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-t', '--time', type=int, default=600,
                        help='time to run in seconds')
    parser.add_argument('-c', '--cleanup', type=int, default=5,
                        help='time to simulate cleanup in seconds')
    return parser.parse_args()


if __name__ == '__main__':
    _logger.info(f'Process running with PID: {os.getpid()}')

    args = parse_args()

    for sig in signal.valid_signals():
        # SIGKILL and SIGSTOP can't be caught
        if not sig in {signal.SIGKILL, signal.SIGSTOP}:
            signal.signal(sig, functools.partial(signal_handler, args.cleanup))

    sys.exit(main(args.time))