Python 装饰器详解

Python 装饰器就是接收一个函数、再返回新函数的函数。我们从零写一个装饰器，接着看看为什么一定要用 functools.wraps。

装饰器是这样一个函数：接收一个函数，返回一个替代品。@ 那一行只是重新给名字赋值的简写。

亲手写过一次之后，这个语法就不再像魔法了。

不用 @ 写一个

内部的 wrapper 接收任意参数，调用原函数，再修改结果。

# a decorator is a function that returns a function
def shout(fn):
    def wrapper(*args, **kwargs):
        return fn(*args, **kwargs).upper()
    return wrapper

def greet(name):
    return f'hello {name}'

loud = shout(greet)
print(loud('ada'))

输出：

HELLO ADA

到这里还没发生什么特别的事。传进去一个函数，出来的是另一个函数。

@ 那一行是一回事

定义上方的 @shout 等同于 greet2 = shout(greet2)。

# the @ syntax is the same thing
@shout
def greet2(name):
    return f'hello {name}'

print(greet2('ada'))

输出：

HELLO ADA

结果一样。装饰器只运行一次，就在定义函数的时候。

包装函数把原函数藏起来了

函数名和文档字符串现在都属于包装函数。

# the wrapper hides the original function
print(greet2.__name__, '|', greet2.__doc__)

输出：

wrapper | None

这会弄坏 help 输出、调试器，以及所有读取 __name__ 的东西。

用 functools.wraps 修好它

@functools.wraps(fn) 会把函数名、文档字符串和其他几个属性复制到包装函数上。

# functools.wraps keeps the metadata
import functools

def shout_fixed(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        return fn(*args, **kwargs).upper()
    return wrapper

@shout_fixed
def greet3(name):
    """Say hello."""
    return f'hello {name}'

print(greet3.__name__, '|', greet3.__doc__)
print(greet3('ada'))

输出：

greet3 | Say hello.
HELLO ADA

你写的每个装饰器都该加上它。没有哪种情况会让你想要包装函数自己的元数据。

带参数的装饰器

@repeat(3) 的意思是“先调用 repeat(3)，再把它的返回值当作装饰器”。这需要再多嵌套一层。

# a decorator that takes arguments needs one more layer
def repeat(times):
    def decorator(fn):
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            return [fn(*args, **kwargs) for _ in range(times)]
        return wrapper
    return decorator

@repeat(3)
def roll():
    return 4

print(roll())

输出：

[4, 4, 4]

一共三层：接收参数的函数、装饰器、包装函数。

保存状态

包装函数是闭包，所以能在多次调用之间保存状态。把状态挂在包装函数上，外面就能读到。

# state in the closure, such as counting calls
def counted(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        wrapper.calls += 1
        return fn(*args, **kwargs)
    wrapper.calls = 0
    return wrapper

@counted
def work(n):
    return n * 2

work(1); work(2); work(3)
print('called', work.calls, 'times')

输出：

called 3 times

标准库里现成的一个

functools.lru_cache 是按参数缓存结果的装饰器。用在朴素的递归斐波那契上，能把指数级的计算量降到线性。

# caching is a decorator in the standard library
@functools.lru_cache(maxsize=None)
def fib(n):
    return n if n < 2 else fib(n - 1) + fib(n - 2)

print(fib(30), fib.cache_info())

输出：

832040 CacheInfo(hits=28, misses=31, maxsize=None, currsize=31)

缓存信息显示了有多少次调用没有执行函数体就直接拿到了答案。

要点

• @decorator 的意思是 name = decorator(name)。
• 装饰器在定义时运行，包装函数在调用时运行。
• 一定要用 functools.wraps，否则函数名和文档字符串就丢了。
• 带参数的装饰器需要三层嵌套函数。

你要写的装饰器大多是日志、计时、缓存、重试或权限检查。它们的结构都和本文第一个例子一样。