Python中的装饰器

装饰器（Decorator）是Python中一个重要部分，它本质上是一个函数，不同于普通函数，装饰器的返回值是一个函数对象。通过利用装饰器，我们可以让其他函数在不做任何代码改动的情况下增加额外的功能，同时也能够让代码更加简洁。

一个简单的装饰器

在开始装饰器之前，我们先写一个简单的普通函数，该函数的作用是统计 n 以内所有质数之和。

def get_prime_sum(n):
    sum = 0
    for i in range(2, n + 1):
        flag = True
        for j in range(2, i):
            if i % j == 0:
                flag = False
                break
        if flag:
            sum += i
    return sum

上面的函数中，当 n 发生变化时，函数返回值及执行耗时是不一样的，如果需要打印日志，如该函数的方法名、注释内容、入参、返回值和执行耗时，那么我们一般可能会这样做：

import time


def get_prime_sum(n):
    """我是 get_prime_sum 的注释"""
    time1 = time.time()
    sum = 0
    for i in range(2, n + 1):
        flag = True
        for j in range(2, i):
            if i % j == 0:
                flag = False
                break
        if flag:
            sum += i
    time2 = time.time()
    print("函数名：{}".format(get_prime_sum.__name__))
    print("函数注释内容：{}".format(get_prime_sum.__doc__))
    print("入参：{}".format(n))
    print("返回值：{}".format(sum))
    print("函数执行耗时：{:.8f}".format(time2 - time1))
    return sum

从结果上来看，这样是没啥问题的。但我们把打印日志的处理直接加到该函数中，可能会有一些不好的地方：

• 如果我们又需要增加些其他小功能，那么就要再次改动该函数，这相当于间接改变了该函数的定义
• 如果我们有很多个不同函数都需要打印日志，那么会导致很多函数都出现重复代码

面对上面这种情况，我们就可以通过 装饰器 来处理。装饰器的架子大概是这样：

def 装饰器名(func):
    def wrapper(*args, **kwargs):
        ...  # 编写额外功能的代码
        return func(*args, **kwargs)
    return wrapper

在装饰器中，需要接收一个函数 func 作为参数，同时它也会返回一个函数 wrapper。在 wrapper(*args, **kwargs) 中，可以接收任意参数，它首先会执行我们编写额外功能的代码，最后再调用原始函数。

如果要使用该装饰器，那么就需要借助Python的 @ 语法，把装饰器置于普通函数的定义处。在上面的例子中，我们使用装饰器后的代码如下：

import time


def log(func):
    def wrapper(*args, **kwargs):
        """我是 wrapper 的注释"""
        time1 = time.time()
        res = func(*args, **kwargs)
        time2 = time.time()
        print("函数名：{}".format(func.__name__))
        print("函数注释内容：{}".format(func.__doc__))
        print("入参 args ：{}".format(args))
        print("入参 kwargs ：{}".format(kwargs))
        print("返回值：{}".format(res))
        print("函数执行耗时：{:.8f}".format(time2 - time1))
        return res
    return wrapper


@log
def get_prime_sum(n):
    """我是 get_prime_sum 的注释"""
    sum = 0
    for i in range(2, n + 1):
        flag = True
        for j in range(2, i):
            if i % j == 0:
                flag = False
                break
        if flag:
            sum += i
    return sum

有了上面这个装饰器，我们可直接应用到多个不同函数中，如果需要再打印额外日志，那么直接修改装饰器函数代码即可。

带参数的装饰器

上面例子中创建的装饰器是不带参数的，如果我们想对装饰器本身传入参数，那么就需要把上面函数改成一个返回装饰器的函数，这时候就需要来个3层嵌套。例如，我们想在使用装饰器时，增加一个 delay 参数，该参数的作用是让程序执行后延迟一段时间，修改装饰器后的代码如下：

import time


def log(delay=0):
    def decorator(func):
        def wrapper(*args, **kwargs):
            """我是 wrapper 的注释"""
            time1 = time.time()
            res = func(*args, **kwargs)
            time.sleep(delay)  # 延迟 delay 秒
            time2 = time.time()
            print("函数名：{}".format(func.__name__))
            print("函数注释内容：{}".format(func.__doc__))
            print("入参 args ：{}".format(args))
            print("入参 kwargs ：{}".format(kwargs))
            print("返回值：{}".format(res))
            print("函数执行耗时：{:.8f}".format(time2 - time1))
            return res
        return wrapper
    return decorator


@log(5.0)
def get_prime_sum(n):
    """我是 get_prime_sum 的注释"""
    sum = 0
    for i in range(2, n + 1):
        flag = True
        for j in range(2, i):
            if i % j == 0:
                flag = False
                break
        if flag:
            sum += i
    return sum

在上面带参数的装饰器中，首先执行 log(5.0) 时会返回一个装饰器 decorator，接着调用 decorator 函数，它的参数是我们的原始函数 func ，最后返回值是 wrapper 函数。

functools.wraps

在上面最开始的例子中，我们把 @log 放到普通函数的定义处，相当于执行了语句：

get_prime_sum = log(get_prime_sum)

在执行过程中，log() 是一个装饰器，其返回的是 wrapper 函数，所以这里的 get_prime_sum 将会指向一个新的函数 wrapper()，这样一来使用装饰器后的函数 get_prime_sum ，它的属性 __name__ 就会发生变化，我们可以打印出来看看：

get_prime_sum = log(get_prime_sum)
print(get_prime_sum.__name__)  # 输出：wrapper
print(get_prime_sum.__doc__)  # 输出：我是 wrapper 的注释

从结果可以看到，当函数 get_prime_sum() 被装饰后，就会从 get_prime_sum 变成 wrapper，这是因为 wrapper 函数重写了我们原始函数的 __name__ 和 __doc__。

对于这种情况，我们可以通过 functools 模块里面的一个装饰器 wraps 来解决，以不带参数的装饰器为例，修改后的代码如下：

import time
import functools


def log(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        """我是 wrapper 的注释"""
        time1 = time.time()
        res = func(*args, **kwargs)
        time2 = time.time()
        print("函数名：{}".format(func.__name__))
        print("入参 args ：{}".format(args))
        print("入参 kwargs ：{}".format(kwargs))
        print("返回值：{}".format(res))
        print("函数执行耗时：{:.8f}".format(time2 - time1))
        return res
    return wrapper


@log
def get_prime_sum(n):
    """我是 get_prime_sum 的注释"""
    sum = 0
    for i in range(2, n + 1):
        flag = True
        for j in range(2, i):
            if i % j == 0:
                flag = False
                break
        if flag:
            sum += i
    return sum

最后再次打印看下效果：

get_prime_sum = log(get_prime_sum)
print(get_prime_sum.__name__)  # 输出：get_prime_sum
print(get_prime_sum.__doc__)  # 输出：我是 get_prime_sum 的注释

总结

在Python中，我们将会在很多地方看到装饰器的应用场景，比如Python中的 @classmethod、@staticmethod、上面例子中使用的 @functools.wraps(func) 等都使用到了装饰器，另外很多第三方库中也会有大量的装饰器应用。

通过使用装饰器，我们不仅让其他函数在不做任何代码改动的情况下增加额外的功能，同时减少了很多重复代码，让代码更加Pythonic。