Source code for common.timeout

import signal



class TimeoutException(Exception):
  pass




class Timeout:
  """
  Timeout context manager.
  For example this code will raise a TimeoutException:
  with Timeout(seconds=5, error_msg="Sleep was too long"):
    time.sleep(10)
  """
  def __init__(self, seconds, error_msg=None):
    if error_msg is None:
      error_msg = f'Timed out after {seconds} seconds'
    self.seconds = seconds
    self.error_msg = error_msg



  def handle_timeout(self, signume, frame):
    raise TimeoutException(self.error_msg)


  def __enter__(self):
    signal.signal(signal.SIGALRM, self.handle_timeout)
    signal.alarm(self.seconds)

  def __exit__(self, exc_type, exc_val, exc_tb):
    signal.alarm(0)