python装饰器举例详解一看就会

写在前面：

此文为视频笔记，视频传送门，感谢B站up主“正月点灯笼”。

一、正常写一段函数

import time

# is_prime函数是判断一个数是不是质数
def is_prime(num):
    if num < 2:
        return False
    elif num == 2:
        return True
    else:
        for i in range(2,num):
            if num%i == 0:
                return False
        return True

# 求2-10000里面哪些数是质数并打印出来
def prime_nums():
    t1 = time.time()
    for i in range(2,10000):
        if is_prime(i): # 如果是质数
            print(i)
    t2 = time.time()
    print(t2-t1)

prime_nums()

输出结果：

...
...
9991
9993
9995
9997
9999
0.012510061264038086

这个我们都能正常理解。
但是由于prime_nums()函数里面包含代码逻辑部分也包含一个计时的模块，可读性较低，要将两者分开。假设很多函数都需要计时，如何复用这个代码？

二、装饰器第一版

import time

# 先写个装饰器
def dispaly_time(func): # 这里的传参是一个函数
    def wrapper():
        t1 = time.time()
        func() # 传进来的函数在这里执行
        t2 = time.time()
        print("time:",t2-t1)
    return wrapper # 最终要把这个返回出去

# is_prime函数是判断一个数是不是质数
def is_prime(num):
    if num < 2:
        return False
    elif num == 2:
        return True
    else:
        for i in range(2,num):
            if num%i == 0:
                return False
        return True

# 求2-10000里面哪些数是质数并打印出来
@dispaly_time
def prime_nums():
    for i in range(2,10000):
        if is_prime(i): # 如果是质数
            print(i)

prime_nums() # 先进装饰器，然后进入自己函数的逻辑

输出结果：

...
...
9991
9993
9995
9997
9999
time: 0.011608123779296875

这样就讲一个函数的逻辑部分和计时部分分开了，且这个装饰器可以一直被大家复用。

但是如果我自己的逻辑函数有返回值怎么办？

三、有返回值的处理办法

import time

# 先写个装饰器
def dispaly_time(func): # 这里的传参是一个函数
    def wrapper():
        t1 = time.time()
        func() # 传进来的函数在这里执行
        t2 = time.time()
        print("time:",t2-t1)
    return wrapper # 最终要把这个返回出去

# is_prime函数是判断一个数是不是质数
def is_prime(num):
    if num < 2:
        return False
    elif num == 2:
        return True
    else:
        for i in range(2,num):
            if num%i == 0:
                return False
        return True

# 统计2-10000的质数个数
@dispaly_time
def prime_nums():
    cnt = 0
    for i in range(2,10000):
        if is_prime(i): # 如果是质数
            cnt += 1
    return cnt

cnt = prime_nums() # 先进装饰器，然后进入自己函数的逻辑
print(cnt)

输出结果：

time: 0.0039861202239990234
None

并没有打印出来
因为return的这个cnt并没有传到装饰器里面，wrapper函数没有返回值，wrapper里面没有接收到这个参数，所以没有取到对应的函数值。
这样写就能够正常做返回：

import time

# 先写个装饰器
def dispaly_time(func): # 这里的传参是一个函数
    def wrapper():
        t1 = time.time()
        result = func() # 传进来的函数在这里执行
        t2 = time.time()
        print("time:",t2-t1)
        return result
    return wrapper # 最终要把这个返回出去

# is_prime函数是判断一个数是不是质数
def is_prime(num):
    if num < 2:
        return False
    elif num == 2:
        return True
    else:
        for i in range(2,num):
            if num%i == 0:
                return False
        return True

# 统计2-10000的质数个数
@dispaly_time
def prime_nums():
    cnt = 0
    for i in range(2,10000):
        if is_prime(i): # 如果是质数
            cnt += 1
    return cnt

cnt = prime_nums() # 先进装饰器，然后进入自己函数的逻辑
print(cnt)

输出结果：

time: 0.3815629482269287
1229

逻辑函数有返回值已经可以处理，但是如果逻辑函数存在参数呢？

四、有参数的处理方式

import time

# 先写个装饰器
def dispaly_time(func): # 这里的传参是一个函数
    def wrapper(*args): # 不知道参数有多少个可以直接写*args
        t1 = time.time()
        result = func(*args) # 传进来的函数在这里执行
        t2 = time.time()
        print("time:",t2-t1)
        return result
    return wrapper # 最终要把这个返回出去

# is_prime函数是判断一个数是不是质数
def is_prime(num):
    if num < 2:
        return False
    elif num == 2:
        return True
    else:
        for i in range(2,num):
            if num%i == 0:
                return False
        return True

# 统计2-10000的质数个数
@dispaly_time
def prime_nums(maxNum):
    cnt = 0
    for i in range(2,maxNum):
        if is_prime(i): # 如果是质数
            cnt += 1
    return cnt

cnt = prime_nums(10000) # 先进装饰器，然后进入自己函数的逻辑
print(cnt)

输出结果：

time: 0.41730189323425293
1229