【转】Python之装饰器

【转】Python之装饰器

本节内容

---

• 必要知识回顾
• 情景模拟
• 装饰器的概念及实现原理
• 回马枪（带参数的装饰器）

一、 必要知识回顾

---

在开始说装饰器之前，需要大家熟悉之前说过的相关知识点：

• 函数即“变量”： 函数名就是一个变量名，它的值就是其对应的函数体；函数体也可以赋值给其它变量，通过这个变量也能调用函数；
• 嵌套函数： 函数内部可以嵌套定义（一层或多层）函数，内部函数可以在函数体内部调用，也可以当做返回值返回；
• 闭包： 在一个嵌套函数中，内部函数可以调用外部非全局变量并且不受外部函数生命周期的影响；
• 高阶函数： 函数的参数可以是函数；

简单来讲，装饰器就是对这些内容的整合和经典应用。如果不了解这些内容，可以查看 这篇文章。

二、情景模拟

---

我们将通过对一个功能需求的分析和解决过程来探究一下“装饰器是什么”以及“装饰器的一些特性”。

1. 场景说明

假设我现在已经定义了一些函数，并且这些函数都已经被线上业务广泛应用。

import time

def func1():
    time.sleep(1)
    print('func1')
    return 'func1'

def func2():
    time.sleep(2)
    print('func2')
    return 'func2'

...

2. 功能需求

现在线上某个业务响应时间过长，需要在不影响线上服务的情况下分别统计这些函数的运行时间来定位故障。

需求解读：

既然不能影响线上业务，那么必然是不能要求函数调用方去更改代码的，这当然包括调用方式。

实现方式：

• 为每个函数单独添加重复的代码
• 把统计运行时间的的代码封装成一个可接收函数作为参数的高阶函数
• 使用嵌套函数改进上面的高阶函数

3. 实现方式及改进过程

初步实现：分别为各个函数添加运行时间统计功能

import time
import sys

def func1():
    start_time = time.time()
    time.sleep(1)
    print('func1')
    end_time = time.time()
    func_name = sys._getframe().f_code.co_name
    print('%s run time is: %s' % (func_name, (end_time - start_time)))
    return 'func1'

def func2():
    start_time = time.time()
    time.sleep(2)
    print('func2')
    end_time = time.time()
    func_name = sys._getframe().f_code.co_name
    print('%s run time is: %s' % (func_name, (end_time - start_time)))
    return 'func2'
...

函数调用：

func1()
func2()

输出结果：

func1
func1 run time is: 1.000861644744873
func2
func2 run time is: 2.0005600452423096

存在的问题：

功能是实现了，但是存在以下几个问题：

• 如果涉及的函数太多，那么需要做大量的重复性工作；
• 如果这些函数分散在不同的模块，且有不同的人维护，那么需要进行协商沟通来保证代码功能一致性；
• 故障解决后，还需要一个一个的去删除或注释调试代码，又是一次大量的重复性工作；
• 如果其他业务也遇到相同问题，需要再次在多个地方复写这些代码。

改进思路：

要避免重复劳动，提高代码重用性，大家很自然就会想到把这个功能封装成一个函数。

改进1：自定义一个（高阶）函数来提高统计代码复用性

import time

def func1():
    time.sleep(1)
    print('func1')
    return 'func1'

def func2():
    time.sleep(2)
    print('func2')
    return 'func2'

def print_run_time(f):
    time_start = time.time()
    ret = f()
    time_end = time.time()
    func_name = f.__name__
    print('%s run time is: %s' % (func_name, (time_end - time_start)))
    return ret

函数调用：

print_run_time(func1)
print_run_time(func2)

输出结果：

func1
func1 run time is: 1.0003290176391602
func2
func2 run time is: 2.0004265308380127

存在的问题：

统计代码的重复性工作解决了，但是新的问题出现了：此时只能通过print_run_time(f)函数去调用原来的函数了（func1, func2, ...），这显然已经改变了函数的调用方式，因此是不合理的。

看上去越改越差劲了，我们继续分析下，会峰回路转的。

改进思路：

如果函数调用方式不能修改，那么我们只能给原来的函数名重新赋值一个新的函数体了，这个新的函数体就应该是添加完统计功能之后的函数体。我们看看下面这两种实现行不行：

func1 = print_run_time
func2 = print_run_time

print_run_time(f)是有参数的，而原函数func1和func2是没参数的，调用方式还是发生改变了，因此这种方式不可行。

func1 = print_run_time(func1)
func2 = print_run_time(func2)

上面这种方式显然更不行了，因为print_run_time(f)返回的是原函数的返回值，而这个返回值不是一个函数，这将导致被重新赋值后的func1和func2无法被调用。

那么我们是否可以把print_run_time(f)函数的函数体定义为一个内部的嵌套函数，然后将这个内部的嵌套函数作为print_run_time(f)函数的返回值呢？这貌似是说的通的，看下面的实现。

改进2：使用嵌套函数和闭包改进上面定义的高阶函数

import time

def func1():
    time.sleep(1)
    print('func1')
    retrun 'func1'

def func2():
    time.sleep(2)
    print('func2')
    return 'func2'

def print_run_time(f):
    def inner():
        time_start = time.time()
        ret = f()
        time_end = time.time()
        func_name = f.__name__
        print('%s run time is: %s' % (func_name, (time_end - time_start)))
        return ret
    return inner
    
func1 = print_run_time(func1)
func2 = print_run_time(func2)

函数调用：

func1()
func2()

输出结果：

func1
func1 run time is: 1.0003256797790527
func2
func2 run time is: 2.000091791152954

Cool! We got it! 我们现在所需要的做的是在所有需要统计运行时长的函数定义之后的任意地方执行一下下面这条语句就可以了：

funcN = print_run_time(funcN)

存在的问题：

我们貌似忽略了一个问题，如果原函数有参数怎么办？

改进思路：

是的，我们定义print_run_time(f)函数的内部函数inner()时，为它定义相应的参数就可以了。由于每个函数的参数数量是不同的，因此inner函数的参数应该是可变(长)参数。

改进3：支持原函数传递参数

def func1(string):
    print(string)
    time.sleep(1)
    print('func1')
    return 'func1 return'

def func2():
    time.sleep(2)
    print('func2')
    
def print_run_time(f):
    def inner(*args, **kwargs):
        time_start = time.time()
        ret = f(*args, **kwargs)
        time_end = time.time()
        func_name = f.__name__
        print('%s run time is: %s' % (func_name, (time_end - time_start)))
        return ret
    return inner

func1 = print_run_time(func1)
func2 = print_run_time(func2)

函数调用：

ret1 = func1('decorator test!')
print(ret1)
ret2 = func2()
print(ret2)

输出结果：

decorator test!
func1
func1 run time is: 1.0001435279846191
func1 return
func2
func2 run time is: 2.0005276203155518
None

到目前为止：统计函数运行时间的功能实现了，函数原来的调用方式没有发生改变，原函数的定义也没有发生改变。其实这就是是装饰器的雏形。

三. 装饰器的概念及原理

---

1. 什么是装饰器？

装饰器，是一种“语法糖”，其本质上就是个函数。

2. 装饰器的作用

它是一个装饰其他函数的函数，用来为其他函数添加一些额外的功能。

3. 装饰器原则

装饰器对被装饰的函数应该是完全透明的，即

• 不能修改被装饰的函数的源代码
• 不能修改被装饰的函数的调用方式

4. 什么样的函数才是装饰器？

高阶函数 + 嵌套函数 => 装饰器

这里的高阶函数需要同时满足以下两个条件：

• 接收函数名作为参数 -- 可以实现在不修改被装饰函数源代码的情况下为其添加新的功能
• 返回内部嵌套函数的函数名 -- 可以实现不用修改函数的调用方式

5. 装饰器实现实例

再来看下上面写的实现代码：

def func1(string):
    print(string)
    time.sleep(1)
    print('func1')
    return 'func1 return'

def func2():
    time.sleep(2)
    print('func2')
    
def print_run_time(f):
    def inner(*args, **kwargs):
        time_start = time.time()
        ret = f(*args, **kwargs)
        time_end = time.time()
        func_name = f.__name__
        print('%s run time is: %s' % (func_name, (time_end - time_start)))
        return ret
    return inner

对于这段代码来讲，print_run_time(f)就已经是一个装饰器函数。为了不改变原函数的调用方式，我们需要把print_run_time(f)函数的返回值重新赋值给原来的函数名：

func1 = print_run_time(func1)
func2 = print_run_time(func2)

但是我们上面说过，装饰器除了是一个函数之外，还是一个“语法糖”。“语法糖”应该是可以简化某些操作的，事实上确实是这样。上面的过程其实可以这样来写：

def print_run_time(f):
    def warpper(*args, **kwargs):
        time_start = time.time()
        ret = f(*args, **kwargs)
        time_end = time.time()
        func_name = f.__name__
        print('%s run time is: %s' % (func_name, (time_end - time_start)))
        return ret
    return wrapper

@print_run_time
def func1(string):
    print(string)
    time.sleep(1)
    print('func1')
    return 'func1 return'

@print_run_time
def func2():
    time.sleep(2)
    print('func2')

是的，就是这么简单。

@print_run_time
def func1(string):
    ...

就相当于：

func1 = print_run_time(func1)

不要问为什么，因为Python解释器就是这样执行的。但是需要注意，此时print_run_time(f)函数必须定义在被修改的函数定义之前，这个很容易理解，只要捋一下代码执行过程就明白了。另外，我们把print_run_time的内部函数名改成了wrapper，这个是装饰器函数的惯用名称（当然，也可以继续使用inner或使用任意名称）。

那么，现在我们可以像原来那样调用函数了（上面的修改这对函数的调用方是完全无感知的）：

func1('Decorator Test...')
func2()

输出结果：

Decorator Test...
func1
func1 run time is: 1.0006635189056396
func2
func2 run time is: 2.0003299713134766

四、 回马枪

---

上面的“情景模拟”部分对于print_run_time函数的改进其实还没有完成。出于内容衔接和便于对概念理解的目的，才将这个改进放到了这里。现在我们返回来，看下这个print_run_time(f)函数还有什么不足。

1. 一个新需求

现在要求在完成统计函数运行时间的基础上，如果函数运行时间超过指定的秒数则打印提示信息，且该秒数允许自定义。

存在的问题：

print_run_time(f)作为一个函数，目前对可而接收的参数限制太大--只能接受一个函数作为参数，要接收一个时间参数就必须为该函数定义新的参数。

解决思路：
这个看起来很容易解决，只要给print_run_time()函数定义一个用于接收超时时间的参数就可以了。

2. 代码初步实现

def print_run_time(f, timeout=1):
    def inner(*args, **kwargs):
        time_start = time.time()
        ret = f(*args, **kwargs)
        time_end = time.time()
        func_name = f.__name__
        run_time = time_end - time_start
        print('%s run time is: %s' % (func_name, run_time))
        if run_time > timeout:
            print('PROBLEM'.rjust(30, '>'))
        return ret
    return inner

def func1(string):
    print(string)
    time.sleep(1)
    print('func1')
    return 'func1 return'

def func2():
    time.sleep(2)
    print('func2')

那么此时，我们可以这样来重新给func1和func2函数赋值并指定超时时间：

func1 = print_run_time(func1, timeout=2)
func2 = print_run_time(func2, timeout=2)

函数调用：

func1('Decorator Test...')
func2()

输出结果：

Decorator Test...
func1
func1 run time is: 1.0006680488586426
func2
func2 run time is: 2.0003504753112793
>>>>>>>>>>>>>>>>>>>>>>>PROBLEM

可以看到，我们设置的超时时间为2秒，func1函数的运行时间为1秒多，小于2秒，因此没有输出错误信息；而func2函数的运行时间为2秒多，大于2秒，因此输出了一个PROBLEM错误信息。Great! It works!

3. 代码改进

来看下这时候“装饰器语法糖”能正常工作吗？

因为print_run_time(f, timeout=1)的两个参数都是位置参数，要指定timeout的值就要先指定f的值，因此写法是这样的：

@print_run_time(func1, timeout=2)
def func1(string):
    ...

很抱歉，这种方式行不通，会报NameError: name 'func1' is not defined。

结合这个例子和前面可以正常执行的代码，我们可以得出一个结论：装饰器函数只能有一个接受函数的参数；或者说以“装饰器语法糖”格式使用的函数，必须本身是一个装饰器函数或者其返回值是一个装饰器函数。

解决思路：

既然f和timeout不能同时存在，但是又必须都存在，那我们就只能在print_run_time函数中定义一个内部函数来把它们分开了。也就是说print_run_time函数需要返回一个定义在它内部的函数，且这个内部函数需要满足一个正常的装饰器函数。那么参数f必然要定义在这个要被返回内部函数中，而参数timeout只能定义在print_run_time()这个外部函数中了。此时print_run_time函数的定义应该是这样的：

import time

def print_run_time(timeout=1):
    def decorator(f):
        def wrapper(*args, **kwargs):
            time_start = time.time()
            ret = f()
            time_end = time.time()
            func_name = f.__name__
            run_time = time_end - time_start
            print('%s run time is: %s' % (func_name, run_time))
            if run_time > timeout:
                print('PROBLEM'.rjust(30, '>'))
        return wrapper
    return decorator

此时的print_run_time已经不再是一个装饰器函数，而是一个返回装饰器函数的函数。func1使用这个函数语法糖应该是这样的：

@print_run_time(timeout=2)
def func1(string):
    print(string)
    time.sleep(1)
    print('func1')
    return 'func1 return'

这个过程是这样的：

decorator_func = print_run_time(timeout=2)

@decortator_func
def func1(string):
    print(string)
    time.sleep(1)
    print('func1')
    return 'func1 return'

函数调用：

func1('Decorator Test...')

输出结果：

Decorator Test...
func1
func1 run time is: 1.0006623268127441

由于timeout是个默认参数（默认值为1），因此也可以不给它传递值，但是那对小括号不能省略：

@print_run_time()
def func1(string):
    print(string)
    time.sleep(1)
    print('func1')
    return 'func1 return'

此时的输出结果中就会多打印一条错误信息：

Decorator Test...
func1
func1 run time is: 1.0006630420684814
>>>>>>>>>>>>>>>>>>>>>>>PROBLEM

4. 最终实现

看上去，上面已经把所有需求都实现了。不要担心，还差一小步而已。此时我们可以打印一下被装饰函数的__name__属性看下：

print(func1.__name__)
print(func2.__name__)

输出结果：

wrapper
wrapper

发现函数func1和func2的名字都变成了wrapper，当然其实还有一些其他属性都变成了print_run_name的内部函数wrapper的属性值。这显然是不合适的，比如一些依赖函数签名的代码就可能会出现错误。解决这个问题也很简单，就是给print_run_name的内部函数warpper应用一个装饰器functools.wrapper就可以了(关于该装饰器的说明请自行翻阅python文档的funcstools模块)：

def print_run_time(timeout=1):
    def decorator(f):
        @functools.wraps(f)
        def wrapper(*args, **kwargs):
            time_start = time.time()
            ret = f(*args, **kwargs)
            time_end = time.time()
            func_name = f.__name__
            run_time = time_end - time_start
            print('%s run time is: %s' % (func_name, run_time))
            if run_time > timeout:
                print('PROBLEM'.rjust(30, '>'))
        return wrapper
    return decorator

此时再打印func1和func2的__name__属性值看看：

print(func1.__name__)
print(func2.__name__)

输出结果：

func1
func2

是的，这次没问题了。

关于Python的装饰器就先说到这里，希望对大家有所帮助。