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 shell(func):
    def run_time():
        "计算参数的函数的运行时间"
        start_time = time.time()
        func()
        stop_time = time.time()
        print(stop_time - start_time)
    return run_time  # 返回这个函数里的run_time子函数的内存地址
def sleep():
    "运行后等待2秒"
    print("wait for 2 seconds")
    time.sleep(2) 
    print("END")
sleep = shell(sleep)  # 赋值之后，sleep现在就是run_time函数的内存地址，也就是附加了新功能的函数
sleep()  # 加上()后就调用执行了sleep现在这个地址里的代码块的内容

run_time现在只是shell函数里的一个子函数，局部变量。也就是调用完shell之后，run_time是会被释放的。

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

sleep=shell(sleep)语句赋值后，sleep不但变为了run_time()函数，而且run_time里的func也已经赋了sleep的内存地址（在调用shell的时候作为参数传入了sleep的内存地址）。

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

最后：

import time
def shell(func):
    def run_time():
        "计算参数的函数的运行时间"
        start_time = time.time()
        func()
        stop_time = time.time()
        print(stop_time - start_time)
    return run_time  # 返回这个函数里的run_time子函数的内存地址
@shell  # 新增了这句，在前面声明，更加易读也更加简洁
def sleep():
    "运行后等待2秒"
    print("wait for 2 seconds")
    time.sleep(2) 
    print("END")
#sleep = shell(sleep)  # 去掉了这句，这句的作用和上面的@shell效果一样。
sleep()

使用@语句，在需要使用装饰器装饰的函数前应用，那么这个函数就会被指定的装饰器装饰。

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

import time
def shell(func):
    def run_time(*args,**kwargs):  # 装饰器的函数使用非固定参数
        "计算参数的函数的运行时间"
        start_time = time.time()
        func(*args,**kwargs)  # 将所有参数导入被装饰的函数
        stop_time = time.time()
        print(stop_time - start_time)
    return run_time  # 返回这个函数里的run_time子函数的内存地址
@shell  # 新增了这句  # 和这句效果一样，sleep=shell(sleep)
def sleep(n):
    "运行后等待n秒"
    print("wait for %d seconds"%n)
    time.sleep(n) 
    print("END")
sleep(3)

如此，参数也不是问题了。还有一个返回值的问题，将sleep修改成一个有返回值的函数。我们又会发现问题了。

另外，再调整一下函数的名字，作为全局变量的函数名需要有意义。把shell改名为run_time，里面的子函数什么名字都无所谓

import time
def run_time(func):
    def wrapper(*args,**kwargs):
        "计算参数的函数的运行时间"
        start_time = time.time()
        res = func(*args,**kwargs)  # 记录被装饰函数的返回值
        stop_time = time.time()
        print(stop_time - start_time)
        return res  # 函数运行结束后，返回被装饰函数的返回值
    return wrapper
@run_time  # 和这句效果一样，sleep=run_time(sleep)
def sleep(n):
    "运行后等待n秒"
    print("wait for %d seconds"%n)
    time.sleep(n) 
    return "END"
print(sleep(2))

最后，在给装饰器也加上参数，实现可以有多种功能的装饰器。我们需要一个三层的装饰器，具体原理也不是很明白，先看三层的代码：

import time
def tmp_wrapper():
    def run_time(func):
        def wrapper(*args,**kwargs):
            "计算参数的函数的运行时间"
            start_time = time.time()
            res = func(*args,**kwargs)
            stop_time = time.time()
            print(stop_time - start_time)
            return res
        return wrapper
    return run_time
@tmp_wrapper()  # 和函数后紧跟的那句赋值效果一样
def sleep(n):
    "运行后等待n秒"
    print("wait for %d seconds"%n)
    time.sleep(n) 
    return "END"
#sleep = tmp_wrapper()(sleep)  # 这句和上面的@的效果一样。
print(sleep(2))

这里@tmp_wrapper() 和sleep=tmp_wrapper()(sleep) 效果一样。写代码的时候使用@，这里我们理解一下赋值的意义。

tmp_wrapper()就是这个函数的返回值，这个返回值就是run_time，也就是run_time的内存地址，后面再跟上(sleep)，就是执行run_time。也就是说sleep=runtime(sleep)。如此剥去了最外面一层，就和上面两层的结构一样了。

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

最后，我们就可以方便的为装饰器带上参数了。我们加上可以控制浮点小数位数的参数，并且把run_time作为最外层的函数名。

三层装饰器模板：

import time
def run_time(n):  # 这里有参数了
    def out_wrapper(func):
        def wrapper(*args,**kwargs):
            "计算参数的函数的运行时间"
            start_time = time.time()
            res = func(*args,**kwargs)
            stop_time = time.time()
            print(round((stop_time - start_time),n))  # 通过装饰器的参数控制小数位数
            return res
        return wrapper
    return out_wrapper
@run_time(8)  # 为装饰器附上参数
def sleep(n):
    "运行后等待n秒"
    print("wait for %d seconds"%n)
    time.sleep(n) 
    return "END"
print(sleep(2))

上面这个就是装饰的最终形态。

再贴一下不带参数的装饰器的模板

两层装饰器模板：

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))

最后这2断代码，就是装饰器的参考模板。使用的时候应该只需要调整一下最外层的函数名和最内层的代码块就可以了，其他部分不用变（也没剩几行了，但主要是结构和逻辑）。就算不理解也能套用了。