"""Simple context managers for handling signals."""

from __future__ import absolute_import
import contextlib
import signal


class SigtermException(Exception):
  """Exception for getting a SIGTERM."""


def _sigterm_handler(signum, frame):
  del signum, frame  # Unused.
  raise SigtermException()


@contextlib.contextmanager
def sigterm_handler():
  """A simple handler for SIGTERM."""

  old_sigterm_handler = signal.signal(signal.SIGTERM, _sigterm_handler)
  try:
    yield
  finally:
    # Reset previous handlers
    signal.signal(signal.SIGTERM, old_sigterm_handler)