timeout.py

import functools
import signal

from contextlib import contextmanager


def signal_handler(_, frame):
    raise TimeoutError(
        'It took too much time to run code in file <{}> on line <{}>'.format(
            frame.f_code.co_filename,
            frame.f_lineno,
        )
    )


def deadline(seconds):
    """
    Deadline decorator that raises an exception if it takes more than :seconds: to run a function

    See more here: https://docs.python.org/3/library/signal.html#example
    
    Note: It will work only for pure python functions

    :param seconds: Number of seconds before the TimeoutError will be raised
    :return:

    Example:
    >>> @deadline(1)
    ... def f():
    ...     while True:
    ...         pass
    ...
    ... try:
    ...     f()
    ... except TimeoutError as e:
    ...     print(e)
    """

    def decorator(f):
        @functools.wraps(f)
        def wrapped(*args, **kwargs):
            # Set the signal handler and a alarm that will be triggered in :seconds: seconds
            signal.signal(signal.SIGALRM, signal_handler)
            signal.alarm(seconds)

            rv = f(*args, **kwargs)

            signal.alarm(0)  # Disable the alarm

            return rv
        return wrapped
    return decorator


@contextmanager
def timeout(seconds):
    """
    Timeout context manager that raises an exception if it takes more than :seconds: to run the code

    See more here: https://docs.python.org/3/library/signal.html#example

    Note: It will work only for pure python functions

    :param seconds: Number of seconds before the TimeoutError will be raised
    :return:

    >>> try:
    ...     with timeout(1):
    ...         while True:
    ...             pass
    ... except TimeoutError as e:
    ...     print(e)
    """
    try:
        # Set the signal handler and a alarm that will be triggered in :seconds: seconds
        signal.signal(signal.SIGALRM, signal_handler)
        signal.alarm(seconds)

        yield
    finally:
        signal.alarm(0)  # Disable the alarm