python - @timeout(timelimit) 装饰器是如何工作的？

Question

我在 Stack Overflow 上发现这个使函数超时的装饰器，我想知道是否有人可以详细解释它是如何工作的，因为代码非常优雅但一点也不清晰。用法是@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 1

How does the @timeout(timelimit) decorator work?

装饰器语法

为了更清楚，基于问题中的示例，用法是这样的:

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

以上是装饰器语法。以下在语义上是等价的:

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

foo = timeout(100)(foo)

请注意，我们将包装原始 foo 的函数命名为“foo”。这就是装饰器语法的含义和作用。

必要的导入

from functools import wraps
import errno
import os
import signal

超时引发异常

class TimeoutError(Exception):
    pass

功能分析

这就是行中的调用，@timeout(timelimit)。 timelimit 参数将被锁定到内部函数中，使这些函数成为“闭包”，之所以这样称呼是因为它们关闭了数据:

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)

这是从原始调用 @timeout(timelimit) 返回装饰器的地方:

    return decorator

包装的好处

wraps 函数允许包装目标函数的函数获取该函数的文档，因为 foo 不再指向原始函数:

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

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

---

更好地使用 wraps

为了进一步说明，wraps 返回一个装饰器，并且打算像这个函数一样使用。最好这样写:

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