装饰器(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。