进阶Python：装饰器 全面详解

Posted 2021-03-25 tmesh-python

进阶Python：装饰器

前言

前段时间我发了一篇讲解Python调试工具PySnooper的文章，在那篇文章开始一部分我简单的介绍了一下装饰器，文章发出之后有几位同学说"终于了解装饰器的用法了"，可见有不少同学对装饰器感兴趣。但是那篇文章主要的目的是在介绍PySnooper，所以没有太深入的展开讲解装饰器，于是在这里就详细的介绍一些装饰器的使用。

装饰器是Python中非常重要的一个概念，如果你会Python的基本语法，你可以写出能够跑通的代码，但是如果你想写出高效、简洁的代码，我认为离不开这些高级用法，当然也包括本文要讲解的装饰器，就如同前面提到的代码调试神器PySnooper一样，它就是主要通过装饰器调用的方式对Python代码进行调试。

什么是Python装饰器？

顾名思义，从字面意思就可以理解，它是用来"装饰"Python的工具，使得代码更具有Python简洁的风格。换句话说，它是一种函数的函数，因为装饰器传入的参数就是一个函数，然后通过实现各种功能来对这个函数的功能进行增强。

为什么用装饰器？

前面提到了，装饰器是通过某种方式来增强函数的功能。当然，我们可以通过很多方式来增强函数的功能，只是装饰器有一个无法替代的优势--简洁。

你只需要在每个函数上方加一个@就可以对这个函数进行增强。

在哪里用装饰器？

装饰器最大的优势是用于解决重复性的操作，其主要使用的场景有如下几个：

• 计算函数运行时间
• 给函数打日志
• 类型检查

当然，如果遇到其他重复操作的场景也可以类比使用装饰器。

简单示例

前面都是文字描述，不管说的怎么天花烂坠，可能都无法体会到它的价值，下面就以一个简单的例子来看一下它的作用。

如果你要对多个函数进行统计运行时间，不使用装饰器会是这样的，

from time import time, sleep

def fun_one():
    start = time()
    sleep(1)
    end = time()
    cost_time = end - start
    print("func one run time {}".format(cost_time))
    
def fun_two():
    start = time()
    sleep(1)
    end = time()
    cost_time = end - start
    print("func two run time {}".format(cost_time))
    
def fun_three():
    start = time()
    sleep(1)
    end = time()
    cost_time = end - start
    print("func three run time {}".format(cost_time))

在每个函数里都需要获取开始时间start、结束时间end、计算耗费时间cost_time、加上一个输出语句。

使用装饰器的方法是这样的，

def run_time(func):
    def wrapper():
        start = time()
        func()                  # 函数在这里运行
        end = time()
        cost_time = end - start
        print("func three run time {}".format(cost_time))
    return wrapper

@run_time
def fun_one():
    sleep(1)
    
@run_time
def fun_two():
    sleep(1)
    
@run_time
def fun_three():
    sleep(1)

通过编写一个统计时间的装饰器run_time，函数的作为装饰器的参数，然后返回一个统计时间的函数wrapper，这就是装饰器的写法，用专业属于来说这叫闭包，简单来说就是函数内嵌套函数。然后再每个函数上面加上@run_time来调用这个装饰器对不同的函数进行统计时间。

可见，统计时间这4行代码是重复的，一个函数需要4行，如果100个函数就需要400行，而使用装饰器，只需要几行代码实现一个装饰器，然后每个函数前面加一句命令即可，如果是100个函数，能少300行左右的代码量。

带参数的装饰器

通过前面简单的例子应该已经明白装饰器的价值和它的简单用法：通过闭包来实现装饰器，函数作为外层函数的传入参数，然后在内层函数中运行、附加功能，随后把内层函数作为结果返回。

除了上述简单的用法还有一些更高级的用法，比如用装饰器进行类型检查、添加带参数的的装饰器等。它们的用法大同小异，关于高级用法，这里以带参数的装饰器为例进行介绍。

不要把问题想的太复杂，带参数的装饰器其实就是在上述基本的装饰器的基础上在外面套一层接收参数的函数，下面通过一个例子说明一下。

以上述例子为基础，前面的简单示例输出的信息是，

func three run time 1.0003271102905273
func three run time 1.0006263256072998
func three run time 1.000312328338623

现在我认为这样的信息太单薄，需要它携带更多的信息，例如函数名称、日志等级等，这时候可以把函数名称和日志等级作为装饰器的参数，下面来时实现以下。

def logger(msg=None):
    def run_time(func):
        def wrapper(*args, **kwargs):
            start = time()
            func()                  # 函数在这里运行
            end = time()
            cost_time = end - start
            print("[{}] func three run time {}".format(msg, cost_time))
        return wrapper
    return run_time

@logger(msg="One")
def fun_one():
    sleep(1)
    
@logger(msg="Two")
def fun_two():
    sleep(1)
    
@logger(msg="Three")
def fun_three():
    sleep(1)
    
fun_one()
fun_two()
fun_three()

可以看出，我在示例基本用法里编写的装饰器外层又嵌套了一层函数用来接收参数msg，这样的话在每个函数(func_one、func_two、func_three)前面调用时可以给装饰器传入参数，这样的输出结果是，

[One] func three run time 1.0013229846954346
[Two] func three run time 1.000720500946045
[Three] func three run time 1.0001459121704102

自定义属性的装饰器

上述介绍的几种用法中其实有一个问题，就是装饰器不够灵活，我们预先定义了装饰器run_time，它就会按照我们定义的流程去工作，只具备这固定的一种功能，当然，我们前面介绍的通过带参数的装饰器让它具备了一定的灵活性，但是依然不够灵活。其实，我们还可以对装饰器添加一些属性，就如同给一个类定义实现不同功能的方法那样。

以输出日志为例，初学Python的同学都习惯用print打印输出信息，其实这不是一个好习惯，当开发商业工程时，你很用意把一些信息暴露给用户。在开发过程中，我更加鼓励使用日志进行输出，通过定义WARNING、DEBUG、INFO等不同等级来控制信息的输出，比如INFO是可以给用户看到的，让用户直到当前程序跑到哪一个阶段了。DEBUG是用于开发人员调试和定位问题时使用。WARING是用于告警和提示。

那么问题来了，如果我们预先定义一个打印日志的装饰器，

def logger_info(func):
    logmsg = func.__name__
    def wrapper():
        func() 
        log.log(logging.INFO, "{} if over.".format(logmsg))
    return wrapper

http://logging.INFO是打印日志的等级，如果我们仅仅写一个基本的日志装饰器logger_info，那么它的灵活度太差了，因为如果我们要输出DEBUG、WARING等级的日志，还需要重新写一个装饰器。

解决这个问题，有两个解决方法：

• 利用前面所讲的带参数装饰器，把日志等级传入装饰器
• 利用自定义属性来修改日志等级

由于第一种已经以统计函数运行时间的方式进行讲解，这里主要讲解第二种方法。

先看一下代码，

import logging
from functools import partial

def wrapper_property(obj, func=None):
    if func is None:
        return partial(attach_wrapper, obj)
    setattr(obj, func.__name__, func)
    return func

def logger_info(level, name=None, message=None):
    def decorate(func):
        
        logmsg = message if message else func.__name__

        def wrapper(*args, **kwargs):
            log.log(level, logmsg)
            return func(*args, **kwargs)

        @wrapper_property(wrapper)
        def set_level(newlevel):
            nonlocal level
            level = newlevel

        @wrapper_property(wrapper)
        def set_message(newmsg):
            nonlocal logmsg
            logmsg = newmsg

        return wrapper

    return decorate


@logger_info(logging.WARNING)
def main(x, y):
    return x + y

这里面最重要的是wrapper_property这个函数，它的功能是把一个函数func编程一个对象obj的属性，然后通过调用wrapper_property，给装饰器添加了两个属性set_message和set_level，分别用于改变输出日志的内容和改变输出日志的等级。

看一下输出结果，

main(3, 3)

# 输出
# WARNING:Test:main
# 6

来改改变一下输出日志等级，

main.set_level(logging.ERROR)
main(5, 5)

# 输出
# ERROR:Test:main
# 10

输出日志等级改成了ERROR。

保留元信息的装饰器

很多教程中都会介绍装饰器，但是大多数都是千篇一律的围绕基本用法在展开，少部分会讲一下带参数的装饰器，但是有一个细节很少有教程提及，那就是保留元信息的装饰器。

什么是函数的元信息？

就是函数携带的一些基本信息，例如函数名、函数文档等，我们可以通过**func.__name__获取函数名、可以通过func.__doc__获取函数的文档信息，用户也可以通过注解**等方式为函数添加元信息。

例如下面代码，

from time import time

def run_time(func):
    def wrapper(*args, **kwargs):
        start = time()
        func()                  # 函数在这里运行
        end = time()
        cost_time = end - start
        print("func three run time {}".format(cost_time))
    return wrapper

@run_time
def fun_one():
    '''
    func one doc.
    '''
    sleep(1)
    
fun_one()

print(fun_one.__name__)
print(fun_one.__doc__)

# 输出
# wrapper
# None

可以看出，通过使用装饰器，函数fun_one的元信息都丢失了，那怎么样才能保留装饰器的元信息呢？

可以通过使用Python自带模块functools中的wraps来保留函数的元信息，

from time import time
from functools import wraps

def run_time(func):
    @wraps(func)                                # <- 这里加 wraps(func) 即可
    def wrapper(*args, **kwargs):
        start = time()
        func()                  # 函数在这里运行
        end = time()
        cost_time = end - start
        print("func three run time {}".format(cost_time))
    return wrapper

@run_time
def fun_one():
    '''
    func one doc.
    '''
    sleep(1)
    
fun_one()

print(fun_one.__name__)
print(fun_one.__doc__)

# 输出
# fun_one   
# func one doc.

只需要在代码中加入箭头所指的一行即可保留函数的元信息。