消除 Python 修饰器的副作用

修饰器

Python 装饰器本质上修改了函数。下面的例子记录了一个函数的运行时间。

import time

def count_time(func):
    def new_func(*args, **kwargs):
        begin = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(func.__name__, 'cost', end-begin, 'seconds')
        return result
    return new_func

@count_time
def is_prime(n: int):
    """Check if n is a prime."""
    if n <= 1:
        return False
    i = 2
    while i * i <= n:
        if n % i == 0:
            return False
        i += 1
    return True

if __name__ == '__main__':
    n = 1000000000 + 7
    p = is_prime(n)
    print(n, 'is prime' if p else 'is not prime')

打印的结果可能是：

is_prime cost 0.00498652458190918 seconds
1000000007 is prime

修饰器函数是修改函数的函数，这是一种高级地抽象。

修饰器的副作用

实际上，修饰器的作用就是将函数修改成修饰器的返回值。这是一种十分简单明了的修饰方法，在多数情况下也是可行的。对于使用修饰器修饰函数的人来说，他的期望是修饰器以修饰器声称的作用修改这个函数，而不用让他感受到其他副作用。

多数情况下，人们对于一个函数的使用，只会通过函数对象的句柄引用这个对象而已。这时，这个函数句柄指向的是原来定义的函数，还是修饰器返回的函数，都没有分别。但是如果对函数的使用超出了对它的引用（或者说用句柄调用），而是要通过句柄访问函数对象的内容，那么修饰器的这个行为就显得比较粗鲁。

接上例：

if __name__ == '__main__':
    n = 1000000000 + 7
    p = is_prime(n)
    print(n, 'is prime' if p else 'is not prime')
    print(is_prime.__name__, is_prime.__doc__)

最后一行打印被修饰函数的名字，和被修饰函数的文档字符串。但是显然我们不能如愿，我们将得到：

new_func None

换言之，得到修饰函数的返回值的名字和文档字符串。

还有一种情况。Python 的对象都可以通过 setattr 来添加子键。这导致函数对象的 __dict__ 信息发生变动。

object.__dict__: A dictionary or other mapping object used to store an object’s (writable) attributes.

我们不仅希望被修饰函数的 __dict__ 不要丢失，而且希望保留修饰器对 __dict__ 的修改。换言之，将修饰器修饰的结果中的 __dict__ 更新。

其三，对于修饰这个行为，应该有一种方法让函数的用户放弃修饰，即访问修饰之前的原函数。函数对象有个 __wrapped__ 成员。所以我们对修饰器的支持应该要设置这个成员，或者至少不能妨碍用户手动设置这个成员。

总结我们的要求：

1. 修改 __name__、__doc__ 等成员。
2. 更新 __dict__。
3. 设置 __wrapper__。

消除副作用

我们期望的是，从被修饰函数中拷贝一些信息到修饰器的返回值里。我们先书写这个过程的抽象。

WRAPPER_ASSIGNMENTS = ('__module__', '__name__', '__doc__', '__qualname__',
                       '__annotations__')
WRAPPER_UPDATES = ('__dict__',)

def update_wrapper(new, old,
                   assigned=WRAPPER_ASSIGNMENTS,
                   updated=WRAPPER_UPDATES):
    for attr in assigned:
        try:
            value = getattr(old, attr)
        except AttributeError:
            pass
        else:
            setattr(new, attr, value)
    for attr in updated:
        getattr(new, attr).update(getattr(old, attr, {}))
    new.__wrapped__ = old
    return new

这个函数提供了对函数进行修复的抽象。任意一个函数的成员要么在修复过程中被赋值，要么像一个 dict 一样调用 update。

Docstring:
D.update([E, ]**F) -> None.  Update D from dict/iterable E and F.
If E is present and has a .keys() method, then does:  for k in E: D[k] = E[k]
If E is present and lacks a .keys() method, then does:  for k, v in E: D[k] = v
In either case, this is followed by: for k in F:  D[k] = F[k]
Type:      method_descriptor

修改修饰器，就可以满足需求：

def count_time(func):
    def new_func(*args, **kwargs):
        begin = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(func.__name__, 'cost', end-begin, 'seconds')
        return result
    return update_wrapper(new_func, func)

注意，如果 update_wrapper 最后并不设置 __wrapped__ 的话。那么想要设置 __wrapped__ 的人来说，他不能在调用 update_wrapper 之前设置 new.__wrapped__ = old，因为 update_wrapper 更新 __dict__ 的部分，可能会覆盖这个赋值。

这个 update_wrapper 实际上就是 functools.update_wrapper。

用修饰器给出消除

update_wrapper 实际上是对修饰器定义的返回函数（上述的 new_func）的修改，所以把他做成修饰器是一个好的实践。

也就是我们需要定义函数 wraps，使得 wraps(old)(new) 等价于 update_wrapper(new, old)。这不难实现：

def wraps(old, assigned=WRAPPER_ASSIGNMENTS, updated=WRAPPER_UPDATES):
    return lambda new: update_wrapper(new, old,
                                      assigned=assigned, updated=updated)

如此一来，我们定义修饰器的过程就得到了简化：

def count_time(func):
    @wraps(func)
    def new_func(*args, **kwargs):
        begin = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(func.__name__, 'cost', end-begin, 'seconds')
        return result
    return new_func

wraps 实际上就是 functools.wraps。只不过标准库不是用 lambda 而是用 partial 实现的。