andry81/tacklelib

# pure python module for commands w/o extension modules usage

import signal, logging

# based on:
#   https://gist.github.com/tcwalther/ae058c64d5d9078a9f333913718bba95
#   http://stackoverflow.com/a/21919644/487556
#
class DelayedSigInterrupt(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, sig = sig):
        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])