装饰器

什么是装饰器？

装饰器的作用

首先要搞清楚装饰器是什么？

$\underline{装饰器是一种“神奇”的函数，它能给被装饰的目标赋予新的能力}$。

装饰器的作用是在不改变原有函数的前提下，为它增加额外的功能。这让我们联想到装饰器设计模式，该模式便是在不改变原有对象的前提下，为其添加新的职责。

我们对装饰器有个基本概念了，但还是有点抽象，不清楚其技术原理。

为了搞懂其原理，我们还需要认识到一点：

函数可以返回函数。

一个装饰器实例

请看实例：

# 吃完想睡函数
def sleep(func):
    def wrapper(food): # !!!wrapper函数的原型和被装饰的函数（如eat）保持一致。
        func(food)
        print("Now, I want some sleep!")
    
    return wrapper
    
def eat(food="fish"):
    print("Eat", food)

eat = sleep(eat)
eat("pork")

运行输出：

Eat pork
Now, I want some sleep!

恭喜你，你已经写了一个装饰器了。

实际上，装饰器就是这么个玩意。

使用@符号

似乎有点不对，我们没有使用“@”符号啊？其实，“@”只是简写形式而已，我们改用@试试就知道了：

$\underline{其实@虽然看起来很神奇，但也没有你想象的那么复杂}$。

# 吃完想睡函数
def sleep(func):
    def wrapper(food):
        func(food)
        print("Now, I want some sleep!")
    
    return wrapper

@sleep  # @sleep等价于上例中的：eat=sleep(eat)
def eat(food="fish"):
    print("Eat", food)

eat("pork")

运行输出：

Eat pork
Now, I want some sleep!

装饰器使用实践：测试代码效率

我们打算使用装饰器测试以下代码的性能：

def generate_list(size=1000000):
    my_list = []
    for num in range(size):
        my_list.append(num)

要测试性能，最笨且最容易被想到的方式是计算开始和结束的时间差。这没有什么意思，不提也罢。

timeit实现

我们也可以使用计时库timeit实现上述目标。

from timeit import timeit

def generate_list(size=1000000):
    my_list = []
    for num in range(size):
        my_list.append(num)
      
    
elapsed = timeit(stmt='generate_list()', setup='from __main__ import generate_list', number=1)

print('Time elapsed: ', elapsed)

运行输出：

Time elapsed:  0.10164559999975609

使用装饰器实现

$\underline{装饰器本质上是一种返回函数的高阶函数}$。我们定义一个能测试函数性能的装饰器，它接受一个函数作为参数，并返回一个函数。由于函数也是一个对象，因此函数对象可以赋值给变量，通过变量就能调用该函数。

import time

def timeit_mine(func):
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        func(*args, **kwargs)
        elasped = time.perf_counter() - start
        print('Time eplased: {}'.format(elasped))
    
    return wrapper

# generate_list被装饰成wrapper函数
# 装饰前的generate_list函数却成为了装饰器timeit_mine的参数(func)，传递给装饰后的wrapper函数。
# 装饰前的generate_list函数的参数，被wrapper接收。
# wrapper通过装饰器接受的参数——实际为装饰器前的generate_list函数，和它所接受的参数，
# 再实现重新调用generate_list的效果，但是在调用前后做点文章，达到装饰的目的。
@timeit_mine
def generate_list(size=1000000):
    my_list = []
    for num in range(size):
        my_list.append(num)
        
print('function generate_list name is {}'.format(generate_list.__name__))
generate_list()
generate_list(10000000)

运行输出：

function generate_list name is wrapper
Time eplased: 0.10633619999862276
Time eplased: 1.0675382999979774

更精细地使用装饰器

经过装饰，generate_list的函数名变成了wrapper。我们还想它保持原名，怎么办？ 如果我们还想装饰器能增加一些测试参数，比如重复次数，怎么办？

参考下例：

import time
import functools

def timeit_mine(repeat=3):
    def decorator(func):
        @functools.wraps(func) # 修正函数名称
        def wrapper(*args, **kwargs):
            for i in range(repeat):
                start = time.perf_counter()
                func(*args, **kwargs)
                elasped = time.perf_counter() - start
                print('Time eplased: {}'.format(elapsed))
        return wrapper
    
    return decorator

@timeit_mine(repeat=2)
def generate_list(size=1000000):
    my_list = []
    for num in range(size):
        my_list.append(num)
        

generate_list()
print('function generate_list name is {}'.format(generate_list.__name__))

运行输出：

Time eplased: 0.10164559999975609
Time eplased: 0.10164559999975609
function generate_list name is generate_list