"""A simple timeout context manager.

Note: this depends on SIGALRM and is thus UNIX specific;
      this will not work on Windows.
"""

from __future__ import absolute_import
import contextlib
import signal


class TimeoutException(Exception):
  """Exception for timing out."""


def _default_timeout_handler(signum, frame):
  del signum, frame  # Unused.
  raise TimeoutException()


@contextlib.contextmanager
def timeout(handler=_default_timeout_handler, timeout_secs=0):
  """A simple timeout contextmanager based on SIGALRM.

  Usage:

  try:
    with timeout(handler=my_handler, timeout_secs=30):
      do_stuff_that_might_take_too_long()
  except timeout.TimeoutException:
    clean_up_from_timeout()
  finally:
    other_clean_up()

  Note: Only one timeout instance can be used at a time, since the
        signal handler is changed and all previously set SIGALARMs
        are cleared.

  Args:
    handler: A function called with two arguments:
               |signum| (the signal number which should be signal.SIGALRM)
               |frame| (the current stack frame, which can be None)
             (Defaults to a function that just raises TimeoutException.)
    timeout_secs: The amount of time before a timeout, in seconds.

  Yields:
    Nothing. Yields is used as part of the contextmanager construct.

  Raises:
    timeout.TimeoutException (if the default handler is used).
  """
  old_handler = signal.signal(signal.SIGALRM, handler)
  old_timeout = signal.alarm(timeout_secs)
  try:
    yield
  except:
    # Catch any exceptions (e.g. TimeoutException) that get raised so that
    # the clean-up in the finally block can happen. Then re-raise the exception.
    raise
  finally:
    signal.signal(signal.SIGALRM, old_handler)  # reset the previous handler
    signal.alarm(old_timeout)  # reset the previous time remaining