timeout.py

"""A context manager for enforcing (integer) timeouts of a code block."""
import contextlib


@contextlib.contextmanager
def Timeout(timeout: Optional[int]) -> ContextManager[None]:
  """Raises TimeoutError() if block isn't completed within `timeout` seconds.

  Args:
    timeout: Seconds. If None, there is no timeout.

  Yields:
    None

  Raises:
    TimeoutError: If timeout is exceeded.
  """
  if timeout is not None and timeout <= 0:
    raise ValueError(f'Invalid timeout {timeout}')

  def _handler(signum, frame):
    logging.debug('Timeout %d reached', timeout)
    raise TimeoutError(frame)

  if timeout:
    signal.signal(signal.SIGALRM, _handler)
    signal.alarm(timeout)

  try:
    yield
  finally:
    if timeout:
      signal.alarm(0)