Home | History | Annotate | Download | only in utils 
      1 # Copyright 2016 The Chromium Authors. All rights reserved.
      2 # Use of this source code is governed by a BSD-style license that can be
      3 # found in the LICENSE file.
      4 
      5 import contextlib
      6 import signal
      7 
      8 
      9 @contextlib.contextmanager
     10 def SignalHandler(signalnum, handler):
     11   """Sets the signal handler for the given signal in the wrapped context.
     12 
     13   Args:
     14     signum: The signal for which a handler should be added.
     15     additional_handler: The handler to add.
     16   """
     17   existing_handler = signal.getsignal(signalnum)
     18 
     19   try:
     20     signal.signal(signalnum, handler)
     21     yield
     22   finally:
     23     signal.signal(signalnum, existing_handler)
     24 
     25 
     26 @contextlib.contextmanager
     27 def AddSignalHandler(signalnum, additional_handler):
     28   """Adds a signal handler for the given signal in the wrapped context.
     29 
     30   This runs the new handler after any existing handler rather than
     31   replacing the existing handler.
     32 
     33   Args:
     34     signum: The signal for which a handler should be added.
     35     additional_handler: The handler to add.
     36   """
     37   existing_handler = signal.getsignal(signalnum)
     38 
     39   def handler(signum, frame):
     40     if callable(existing_handler):
     41       existing_handler(signum, frame)
     42     additional_handler(signum, frame)
     43 
     44   try:
     45     signal.signal(signalnum, handler)
     46     yield
     47   finally:
     48     signal.signal(signalnum, existing_handler)
     49