Python自动化开发学习4-装饰器

Posted 2020-10-13

装饰器

通过装饰器可以为函数添加附加功能。通过修改函数本身也可以实现增加功能，但是通过装饰器还有下面2个好处。装饰器的2个原则：

1. 不改变被装饰函数的源代码

2. 不改变被装饰函数的调用方式
   

遵循了上面2个原则，我们可以在为函数添加附加功能的时候，不必去破坏已有的稳定的代码和代码调用方式。并且易于回退。

先定义一个函数：

import time
def sleep():
    "运行后等待2秒"
    print("wait for 2 seconds")
    time.sleep(2) 
    print("END")
sleep()

现在我希望为这个函数附加一个功能，计算这个函数运行的时间，但是我不希望修改sleep函数的源代码，所以我们再定义一个函数来附加这个功能

import time
def sub_a(func):
    "计算参数的函数的运行时间"
    t = time.time()  # 记录当前时间戳
    func()
    print(time.time() - t)  # 现在的时间戳和之前的时间戳的差值就是中间这个函数的运行时间
def sleep():
    "运行后等待2秒"
    print("wait for 2 seconds")
    time.sleep(2) 
    print("END")
sub_a(sleep)

我们未修改sleep函数，而是定义了一个新函数。通过这个新的函数执行附加的代码并且用参数调用执行被装饰的函数。

但是这里我们的调用方式变了，现在需要调用新的函数，并且要使用原来的调用方式sleep()来调用。

在之前的基础上，我们给装饰器外面再套上一层函数。先看代码，下面再分析实现过程。

import time
def sub_b(func):
    def sub_a():
        "计算参数的函数的运行时间"
        t = time.time()
        func()
        print(time.time() - t)
    return sub_a  # 返回sub_a函数的内存地址
def sleep():
    "运行后等待2秒"
    print("wait for 2 seconds")
    time.sleep(2) 
    print("END")
sleep = sub_b(sleep)
sleep()

sleep = sub_b(sleep)语句赋值后，sleep的值是sub_b函数的返回值，也就是sub_a的内存地址，而且sub_a里的func也已经赋了sleep的内存地址（在调用sub_b的时候作为参数传入了sleep的内存地址）。

最后我们再sleep()实际调用的是sub_a()，而sub_a里的func()执行的是sleep()。附加代码和sleep的源代码都执行了，并且最后的调用方式也没变。

这里再插入解释一个小问题，sub_a现在只是sub_b函数里的一个子函数，局部变量。也就是调用完sub_b之后，sub_a是会被释放的。

但是由于sub_b最后的return sub_a返回了sub_a的内存地址，并且在sleep=sub_b(sleep)，将这个内存地址又赋值给了sleep。所以在sub_b调用结束后，sub_a还再被sleep占用着，所以并没有释放sub_a。

至此不改变调用方式的原则也实现了。

再看看参数，之前使用的函数都没有参数，如果为sleep()函数加一个参数，那么代码就报错。所以目前的装饰器只能装饰不带参数的函数，要适应各种参数类型的函数，我们需要使用不固定参数

import time
def sub_b(func):
    def sub_a(*args,**kwargs):  # 这里无脑的加上不固定参数
        "计算参数的函数的运行时间"
        t = time.time()
        func(*args,**kwargs)  # 将所有参数都传给被装饰函数
        print(time.time() - t)
    return sub_a
def sleep(n):  # 现在我们换一个有参数的函数来装饰它
    "运行后等待n秒"
    print("wait for %d seconds"%n)
    time.sleep(n) 
    print("END")
sleep = sub_b(sleep)
sleep(2)

简单无脑的加上不固定参数就行了，具体是什么参数，具体有多少参数我们不用关心，只要全部传过去就行了。

如此，参数也不是问题了。还有一个返回值的问题，将sleep修改成一个有返回值的函数。又会有问题，这次不报错，但是返回值是None。

再修改，调用被装饰函数时，将返回值保存。装饰器最后再return这个值给装饰器。

import time
def sub_b(func):
    def sub_a(*args,**kwargs):
        "计算参数的函数的运行时间"
        t = time.time()
        res = func(*args,**kwargs)  # 将被装饰函数的返回值保存
        print(time.time() - t)
        return res  # 将之前保存的值在这里返回，作为装饰器的返回值
    return sub_a
def sleep(n):
    "运行后等待n秒"
    print("wait for %d seconds"%n)
    time.sleep(n) 
    return "END"
sleep = sub_b(sleep)
print(sleep(2))

现在返回值也有了。东西都全了，但是和别人的装饰器还有一点点不一样。也就差了两句。别人的装饰器大概这下面这个样子的。

import time
def sub_b(func):
    def sub_a(*args,**kwargs):
        "计算参数的函数的运行时间"
        t = time.time()
        res = func(*args,**kwargs)
        print(time.time() - t)
        return res
    return sub_a
@sub_b  # 添加这一句
def sleep(n):
    "运行后等待n秒"
    print("wait for %d seconds"%n)
    time.sleep(n) 
    return "END"
#sleep = sub_b(sleep)  # 去掉这一句
print(sleep(2))

这里使用了@语句，并且放到了被装饰函数前面。替代了原本在被装饰函数后面的赋值语句。这两句效果是一样的。

但是@语句明显代码更加简洁，并且位置是在被装饰函数前，更加易于阅读。

不过后面还会继续使用赋值的形式，我觉得更好理解装饰器的实现方式。

装饰器实际也是一个函数，那么函数就可以带参数，通过参数，可以使函数有更多的选项。现在我们要给装饰器也加上参数。

之前实现不改变调用方式的时候，我们为装饰器套了一层。现在为了让装饰器能够带上参数，我们需要再套一层，具体原理也不是很明白，先看三层的代码：

import time
def sub_c():  # 先不写参数，将来参数可以填这里了
    def sub_b(func):
        def sub_a(*args,**kwargs):
            "计算参数的函数的运行时间"
            t = time.time()
            res = func(*args,**kwargs)
            print(time.time() - t)
            return res
        return sub_a
    return sub_b
def sleep(n):
    "运行后等待n秒"
    print("wait for %d seconds"%n)
    time.sleep(n) 
    return "END"
sleep = sub_c()(sleep)
print(sleep(2))

先看一下赋值的意义，sleep=sub_c()(sleep)。sub_c()的值就是他的返回值sub_b，也就是sub_b函数的内存地址。后面再跟上(sleep)，就是执行sub_b(sleep)。也就是说sleep=sub_b(sleep)。如此剥去了最外面一层，就和上面两层的赋值一样了。

所以三层装饰器用的时候，sub_c后面多了一个()，因为这里我们要的是最外层函数的返回值。而二层装饰器用的时候，后面是没有()的，这里要的是函数代码块的内存地址。

最后，我们就可以方便的为装饰器带上参数了。我们加上可以控制浮点小数位数的参数。

import time
def sub_c(n):
    def sub_b(func):
        def sub_a(*args,**kwargs):
            "计算参数的函数的运行时间"
            t = time.time()
            res = func(*args,**kwargs)
            print(round((time.time() - t),n))
            return res
        return sub_a
    return sub_b
def sleep(n):
    "运行后等待n秒"
    print("wait for %d seconds"%n)
    time.sleep(n) 
    return "END"
sleep = sub_c(8)(sleep)
print(sleep(2))

最后还是要用@的形式

总结

最后是2个装饰器最终的模板，按下面说的调整一下，就算不理解应该也能套用了。

1. 将最外层的函数名改成你的装饰器的名字，为你的装饰器起一个合适的名字

2. 内层的函数名无所谓

3. 最内层的代码块做替换成你自己的代码，res=func(*args,**kwargs)和最后的return res这两句不用修改

两层装饰器模板：

import time
def run_time(func):  # 最外层的函数名字换成你装饰器的名字
    def wrapper(*args,**kwargs):
        "计算参数的函数的运行时间"
        t = time.time()  # 这里替换成你的代码段
        res = func(*args,**kwargs)  # 这句不变
        print(time.time() - t)  # 这里换成你的代码段
        return res  # 这句不变
    return wrapper
@run_time
def sleep(n):
    "运行后等待n秒"
    print("wait for %d seconds"%n)
    time.sleep(n) 
    return "END"
print(sleep(2))

三层装饰器模板：

import time
def run_time(n):  # 最外层的函数名字换成你装饰器的名字
    def decorator(func):  # 内层的函数名字无所谓
        def wrapper(*args,**kwargs):
            "计算参数的函数的运行时间"
            t = time.time()  # 这里替换成你的代码段
            res = func(*args,**kwargs)  # 这句不变
            print(round((time.time() - t),n))  # 这里替换成你的代码段
            return res  # 这句不变
        return wrapper
    return decorator
@run_time(8)
def sleep(n):
    "运行后等待n秒"
    print("wait for %d seconds"%n)
    time.sleep(n) 
    return "END"
print(sleep(2))