sakurashinken sakurashinken - 3 months ago 11
Python Question

How does the @timeout(timelimit) decorator work?

I found this decorator that times out a function here on Stack Overflow, and I am wondering if someone could explain in detail how it works, as the code is very elegant but not clear at all. Usage is

@timeout(timelimit)
.

from functools import wraps
import errno
import os
import signal

class TimeoutError(Exception):
pass

def timeout(seconds=100, error_message=os.strerror(errno.ETIME)):
def decorator(func):
def _handle_timeout(signum, frame):
raise TimeoutError(error_message)

def wrapper(*args, **kwargs):
signal.signal(signal.SIGALRM, _handle_timeout)
signal.alarm(seconds)
try:
result = func(*args, **kwargs)
finally:
signal.alarm(0)
return result

return wraps(func)(wrapper)

return decorator

Answer

How does the @timeout(timelimit) decorator work?

Decorator Syntax

To be more clear, the usage is like this:

@timeout(100)
def foo(arg1, kwarg1=None):
    '''time this out!'''
    something_worth_timing_out()

The above is the decorator syntax. The below is exactly equivalent:

def foo(arg1, kwarg1=None):
    '''time this out!'''
    something_worth_timing_out()

foo = timeout(100)(foo)

Note that we name the function that wraps foo, "foo". That's what the decorator syntax means and does.

Necessary imports

from functools import wraps
import errno
import os
import signal

Exception to raise on Timeout

class TimeoutError(Exception):
    pass

Analysis of the function

This is what's called in the line, @timeout(timelimit). These arguments will be locked into the underlying functions, making those functions "closures" so-called because they close-over the data:

def timeout(seconds=100, error_message=os.strerror(errno.ETIME)):

This will return a function that takes a function as an argument, which the next line proceeds to define. This function will return a function that wraps the original function. :

    def decorator(func):

This is a function to timeout the decorated function:

        def _handle_timeout(signum, frame):
            raise TimeoutError(error_message)

And this is the actual wrapper. Before calling the wrapped function, it sets a signal that will interrupt the function if it does not finish in time with an exception:

        def wrapper(*args, **kwargs):
            signal.signal(signal.SIGALRM, _handle_timeout)
            signal.alarm(seconds)
            try:
                result = func(*args, **kwargs)
            finally:
                signal.alarm(0)

This will return the result if the function completes:

            return result

This returns the wrapper. It makes sure the wrapped function gets the attributes from the original function, like docstrings, name, function signature...

        return wraps(func)(wrapper)

and this is where the decorator is returned, from the original call, @timeout(timelimit):

    return decorator

Benefit of wraps

The wraps function allows the function that wraps the target function to get the documentation of that function, because foo no longer points at the original function:

>>> help(foo)
Help on function foo in module __main__:

foo(arg1, kwarg1=None)
    time this out!

Better usage of wraps

To further clarify, wraps returns a decorator, and is intended to be used much like this function. It would be better written like this:

def timeout(seconds=100, error_message=os.strerror(errno.ETIME)):
    def decorator(func):
        def _handle_timeout(signum, frame):
            raise TimeoutError(error_message)
        @wraps(func)
        def wrapper(*args, **kwargs):
            signal.signal(signal.SIGALRM, _handle_timeout)
            signal.alarm(seconds)
            try:
                result = func(*args, **kwargs)
            finally:
                signal.alarm(0)
            return result
        return wrapper
    return decorator