当前位置：首页 > article >正文

【python】装饰器

article 2024/12/25 6:01:39

在Python中，装饰器（decorators）是一种高级功能，它允许你在不修改原有函数或方法代码的情况下，给函数或方法添加新的功能。装饰器本质上是一个返回函数的函数，它接受一个函数作为参数，并返回一个新的函数（有时是对原函数的封装或修改）。

装饰器的基本语法

一个装饰器的定义通常如下：

# 定义装饰器
def my_decorator(func):
    def wrapper(*args, **kwargs):
        # 在这里添加你想在函数执行前执行的代码
        print("Something is happening before the function is called.")
        
        # 调用原函数并获取结果
        result = func(*args, **kwargs)
        
        # 在这里添加你想在函数执行后执行的代码
        print("Something is happening after the function is called.")
        
        # 返回原函数的结果
        return result
    
    # 返回wrapper函数，它现在会替换原来的函数
    return wrapper

# 使用装饰器
@my_decorator
def target_function(arg1, arg2):
    pass  # 原始函数的实现

语法说明：

decorator_function：是一个装饰器函数，它接受一个函数func作为参数，并返回一个内部函数wrapper，在wrapper函数内部，你可以执行一些额外的操作，然后调用原始函数func，并返回其结果。
wrapper：是内部函数，它是实际会被调用的新函数，它包裹了原始函数的调用，并在其前后增加了额外的行为。
当我们使用@my_decorator装饰在target_function函数定义前，Python会自动将target_function函数作为参数传递给my_decorator函数，然后将返回的wrapper函数替换掉原来的target_function函数。

装饰器的使用

import time

def log(func):
    def wrapper(*args, **kwargs):
        print("before ", time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), args, kwargs)
        result = func(*args, **kwargs)
        print("after ", time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), result)
        return result

    return wrapper

@log
def add(a: int, b: int) -> int:
    return a + b

add(3, 4)

运行结果如下：

before  2024-12-18 10:33:58 (3, 4) {}
after  2024-12-18 10:33:58 7

当你调用add()函数时，实际上调用的是wrapper函数，它会在原函数add()执行前后添加额外的功能。

带参数的装饰器

def repeat(n:int):
    def decorator(func):
        def wrapper(*args, **kwargs):
            for _ in range(n):
                func(*args, **kwargs)
        return wrapper
    return decorator

@repeat(3)
def hello(name):
    print("hello ", name)

hello("Morris")

运行结果如下：

hello  Morris
hello  Morris
hello  Morris

repeat()函数是一个带参数的装饰器，它接受一个整数参数n，然后返回一个装饰器函数。该装饰器函数内部定义了wrapper函数，在调用原始函数之前重复执行n次。因此，hello()函数在被@repeat(3)装饰后，会打印三次问候语。

保留函数元数据

使用装饰器时，原函数的元数据（如名称、文档字符串等）会被wrapper函数的元数据所覆盖。

# 上面hello()函数的名称已被wrapper函数覆盖
print(hello.__name__) # wrapper

为了保留原函数的元数据，你可以使用functools.wraps装饰器：

from functools import wraps

def repeat2(n:int):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for _ in range(n):
                func(*args, **kwargs)
        return wrapper
    return decorator

@repeat2(3)
def hello2(name):
    print("hello ", name)

print(hello2.__name__) # hello2

类装饰器

除了函数装饰器，Python还支持类装饰器。类装饰器是包含__call__方法的类，它接受一个函数作为参数，并返回一个新的函数。

class DecoratorClass:
    def __init__(self, func):
        self.func = func
    
    def __call__(self, *args, **kwargs):
        # 在调用原始函数之前/之后执行的代码
        result = self.func(*args, **kwargs)
        # 在调用原始函数之后执行的代码
        return result

@DecoratorClass
def my_function():
    pass