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
5import contextlib
6import signal
7
8
9@contextlib.contextmanager
10def 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
27def 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