Python 装饰器


Python装饰器是什么

就来看上面的例子,在性能测试的时候,Locust会收集每次执行task的时间,如果我们平时想要测试函数的执行时间需要怎么办呢:

比如用加法来举例:

import time

def add(a, b):
    start_time = time.time()
    res = a + b
    exec_time = time.time() - start_time
    print("add函数,花费的时间是:{}".format(exec_time))
    return res

那如果我们想复用测试时间的函数呢?

import time

def time_calc_self(func):
    start_time = time.time()
    f = func()
    exec_time = time.time() - start_time
    print(f"{func.__name__}执行时间为{exec_time}秒")
    return f

这样的话func()就只能执行无参数的函数,那如果需要传入参数呢?

万能的python给了我们另一个解决方法:闭包

def outer(x):
    def inner(y):
        return x + y
    return inner

print(outer(6)(5))

像这样使用函数的嵌套定义我们就可以同时传入两层函数的参数进行执行,由此我们联想到,是不是我们用两层嵌套,第一层用来传入函数,第二层用来传入参数就可以啦,于是我们可以定义这样一个函数:

def time_calc(func):
    def wrapper(*args, **kargs):        
        start_time = time.time()        
        f = func(*args,**kargs)        
        exec_time = time.time() - start_time        
        print(f"{func.__name__}执行时间为{exec_time}秒")
        return f    
    return wrapper 

再像上面一样调用:

print(time_calc(add)(1, 2))

就可以同时得到时间和输出啦!

重点来了!!!

那说了这么多,装饰器是什么呢,其实就是一个语法糖,用来代替上面这样的调用

我们只需要在定义函数的时候使用@装饰器名字就可以实现上面的功能啦:

import time

# 定义装饰器
def time_calc(func):
    def wrapper(*args, **kargs):        
        start_time = time.time()        
        f = func(*args,**kargs)        
        exec_time = time.time() - start_time     
        print(f"{func.__name__}执行时间为{exec_time}秒")
        return f    
    return wrapper   

# 使用装饰器
@time_calc    
def add(a, b):
    return a + b
    
@time_calc
def sub(a, b):    
    return a - b

print(add(1, 2))
print(sub(1, 2))


装饰器的作用

装饰器的主要作用是**在不修改原函数代码的情况下,增加函数的新功能**。比如上面的例子,我们给普通的加法和减法函数增加了计时的功能,而不需要在每个函数中都手动添加计时代码。

装饰器不仅可以用来计时,还可以用来**日志记录、权限校验、缓存结果**等等。比如我们想记录一个函数的调用次数,就可以这样写:

def log_call_times(func):
    call_count = 0
    def wrapper(*args, **kwargs):
        nonlocal call_count
        call_count += 1
        print(f"{func.__name__}被调用了{call_count}次")
        return func(*args, **kwargs)
    return wrapper

@log_call_times
def say_hello():
    print("Hello!")

say_hello()
say_hello()


装饰器的原理

装饰器本质上是一个**返回函数的高阶函数**。它接收一个函数作为参数,返回一个新的函数,这个新函数在执行原函数的基础上,增加了额外的功能。

在上面的例子中,`time_calc` 就是一个装饰器。当我们写 `@time_calc` 的时候,Python 会自动将 `add` 函数作为参数传递给 `time_calc`,然后将 `time_calc` 返回的 `wrapper` 函数赋值给 `add`。所以当我们调用 `add(1, 2)` 的时候,实际上是在调用 `wrapper(1, 2)`。

带参数的装饰器

有时候我们希望装饰器能够接收参数,比如我们想给计时装饰器增加一个阈值,当函数执行时间超过这个阈值时才打印警告信息。这时候就需要定义一个**带参数的装饰器**:

def time_calc_with_threshold(threshold):
    def decorator(func):
        def wrapper(*args, **kwargs):
            start_time = time.time()
            result = func(*args, **kwargs)
            exec_time = time.time() - start_time
            if exec_time > threshold:
                print(f"警告:{func.__name__}执行时间超过了{threshold}秒,实际执行时间为{exec_time}秒")
            return result
        return wrapper
    return decorator

@time_calc_with_threshold(threshold=0.5)
def slow_function():
    time.sleep(1)

slow_function()

在这个例子中,`time_calc_with_threshold` 是一个带参数的装饰器。它接收一个参数 `threshold`,然后返回一个真正的装饰器 `decorator`。`decorator` 再接收一个函数 `func`,返回一个新的函数 `wrapper`。

装饰器的注意事项

- **装饰器会改变函数的元信息**:比如函数名、文档字符串等。如果想保留原函数的元信息,可以使用 `functools.wraps` 装饰器。例如:

import functools

def time_calc(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        exec_time = time.time() - start_time
        print(f"{func.__name__}执行时间为{exec_time}秒")
        return result
    return wrapper

- **装饰器的执行顺序**:如果一个函数被多个装饰器装饰,装饰器的执行顺序是从**离函数最近的装饰器开始**。例如:

@decorator1
@decorator2
def func():
    pass

这个例子中,`decorator2` 会先被应用到 `func` 上,然后再将结果传递给 `decorator1`。

装饰器是 Python 中一个非常强大且灵活的特性,掌握它可以让代码更加简洁、优雅,同时也能提高代码的复用性和可维护性。希望这篇文章能帮助你更好地理解装饰器!