Timeout decorator and context manager using signals that allows nested timeouts

timeout.py

from contextlib import contextmanager
import errno
import functools
import os
import signal


class timeout:
    _disabled = 0

    def __init__(self, seconds=60, error_message=os.strerror(errno.ETIMEDOUT)):
        self.seconds = seconds
        self.error_message = error_message

    def handle_timeout(self, signum, frame):
        raise TimeoutError(self.error_message)

    @classmethod
    def __setup(cls, seconds, handler):
        # keep track of the current handler that was set
        updated = False
        cur_handler = signal.getsignal(signal.SIGALRM)
        cur_timeout, _ = signal.getitimer(signal.ITIMER_REAL)
        # only set timer if it is not set or is shorter than the current timer
        if seconds < cur_timeout or cur_timeout == 0.0:
            signal.signal(signal.SIGALRM, handler)
            signal.setitimer(signal.ITIMER_REAL, seconds)
            updated = True
        return updated, cur_handler, cur_timeout

    @classmethod
    def __teardown(cls, seconds, updated, handler, timeout):
        signal.signal(signal.SIGALRM, handler)
        if not updated:
            # the original timer was left unchanged, do nothing
            pass
        elif timeout == 0.0:
            # there was no active timer, turn it off
            signal.setitimer(signal.ITIMER_REAL, 0)
        else:
            # a longer timer was in place, return its value minus elapsed time
            inner_time, _ = signal.getitimer(signal.ITIMER_REAL)
            elapsed = seconds - inner_time
            # if we were close to outer timeout make sure we dont get a negative (or zero)
            outer_time = max(0.001, timeout - elapsed)
            signal.setitimer(signal.ITIMER_REAL, outer_time)

    def __enter__(self):
        if self._disabled:
            return
        self._state = self.__setup(self.seconds, self.handle_timeout)

    def __exit__(self, type, value, traceback):
        if self._disabled:
            return
        self.__teardown(self.seconds, *self._state)

    @classmethod
    @contextmanager
    def disable(cls):
        cls._disabled += 1
        try:
            yield
        finally:
            cls._disabled -= 1

    @classmethod
    def timeout(cls, seconds=60, error_message=os.strerror(errno.ETIMEDOUT)):
        def decorator(func):
            def _handle_timeout(signum, frame):
                raise TimeoutError(error_message)

            @functools.wraps(func)
            def wrapper(*args, **kwargs):
                if not cls._disabled:
                    state = cls.__setup(seconds, _handle_timeout)
                try:
                    result = func(*args, **kwargs)
                finally:
                    if not cls._disabled:
                        cls.__teardown(seconds, *state)
                return result

            return wrapper

        return decorator


### DEMO ###

import time


@timeout.timeout(seconds=2, error_message="FUNC")
def do_work(n):
    time.sleep(n)


# change numbers around to see different timeouts triggered
with timeout(seconds=5, error_message="OUTER"):
    time.sleep(1)
    with timeout(seconds=0.5, error_message="INNER"):
        do_work(1)
    time.sleep(5)

# disable timeouts inside a context
with timeout(5, "OUTER"):
    with timeout.disable():
        do_work(3)
        with timeout(0.1):
            time.sleep(1)
        do_work(3)