python 装饰器

装饰器是 Python 中的一种设计模式，用于在不修改函数或类的源代码的情况下，动态地增加或扩展其功能。装饰器本质上是一个高阶函数，它接收一个函数或类作为参数，并返回一个新的函数或类。

装饰器的基本语法

装饰器通常使用 @decorator_name 的语法糖来应用。在函数定义之前加上 @decorator_name，Python 会将该装饰器应用到函数上。

@decorator_name
def function_to_decorate():
    pass

等价于：

def function_to_decorate():
    pass

function_to_decorate = decorator_name(function_to_decorate)

简单装饰器示例

1. 无参数的函数装饰器

一个简单的装饰器，用于在函数执行前后打印消息：

def my_decorator(func):
    def wrapper():
        print("函数执行之前")
        func()
        print("函数执行之后")
    return wrapper

@my_decorator
def say_hello():
    print("Hello, World!")

say_hello()

输出：

函数执行之前
Hello, World!
函数执行之后

在这个例子中，my_decorator 是一个装饰器，它将 say_hello 函数包装在 wrapper 函数中，在函数执行前后打印消息。

2. 带参数的函数装饰器

如果装饰的函数接受参数，装饰器的 wrapper 函数也需要接受这些参数并传递给原始函数：

def my_decorator(func):
    def wrapper(*args, **kwargs):
        print("函数执行之前")
        func(*args, **kwargs)
        print("函数执行之后")
    return wrapper


@my_decorator
def add(a, b):
    print(a + b)


add(5, 3)

输出：

函数执行之前
8
函数执行之后

在这个例子中，装饰器 my_decorator 能够处理任意数量的参数，并且在函数执行前后打印消息。

3.有返回值的函数装饰器

如果装饰的函数有返回值，装饰器的 wrapper 函数需要接收原始函数的返回值并作为其返回值返回：

def my_decorator(func):
    def wrapper(*args, **kwargs):
        print("函数执行之前")
        result = func(*args, **kwargs)
        print("函数执行之后")
        return result
    return wrapper


@my_decorator
def add(a, b):
    return a+b


print(add(5, 3))

输出：

函数执行之前
函数执行之后
8

装饰器链

可以为同一个函数应用多个装饰器，装饰器将按从上到下的顺序依次应用：

def decorator1(func):
    def wrapper(*args, **kwargs):
        print("装饰器1")
        result = func(*args, **kwargs)
        return result
    return wrapper

def decorator2(func):
    def wrapper(*args, **kwargs):
        print("装饰器2")
        result = func(*args, **kwargs)
        return result
    return wrapper

@decorator1
@decorator2
def say_hello():
    print("Hello")

say_hello()

输出：

代码装饰器1
装饰器2
Hello

这里，say_hello 函数首先被 decorator2 装饰，然后再被 decorator1 装饰。执行时，decorator1 先执行，再调用 decorator2。

带参数的装饰器

装饰器本身可以接受参数，使用这种形式的装饰器可以自定义行为。为了实现带参数的装饰器，需要多一层函数嵌套：

def repeat(num_times):
    def decorator(func):
        def wrapper(*args, **kwargs):
            for _ in range(num_times):
                result = func(*args, **kwargs)
            return result
        return wrapper
    return decorator

@repeat(num_times=3)
def greet(name):
    print(f"Hello, {name}")

greet("Alice")

输出：

Hello, Alice
Hello, Alice
Hello, Alice

在这个例子中，repeat(num_times=3) 是一个带参数的装饰器，greet 函数将被调用 3 次。

类装饰器

除了函数装饰器之外，你还可以创建类装饰器。类装饰器是通过定义 __call__() 方法的类来实现的，这样类的实例就可以像函数一样调用：

class MyDecorator:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        print("类装饰器：函数执行之前")
        result = self.func(*args, **kwargs)
        print("类装饰器：函数执行之后")
        return result

@MyDecorator
def say_goodbye(name):
    print(f"Goodbye, {name}")

say_goodbye("Alice")

输出：

类装饰器：函数执行之前
Goodbye, Alice
类装饰器：函数执行之后

在这个例子中，MyDecorator 类实现了一个类装饰器，它在 say_goodbye 函数执行前后打印消息。

使用 functools.wraps

当你使用装饰器时，原始函数的元数据（如函数名和文档字符串）可能会丢失。为了保留这些元数据，可以使用 functools.wraps：

import functools

def my_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print("函数执行之前")
        return func(*args, **kwargs)
    return wrapper

@my_decorator
def example():
    """这是一个示例函数"""
    print("函数正在运行")

print(example.__name__)  # 输出: example
print(example.__doc__)   # 输出: 这是一个示例函数

使用 functools.wraps 后，装饰器不会覆盖原始函数的 __name__ 和 __doc__ 等属性。

实战

1. 编写修饰器，记录函数运行时间长

def runtime(fun):
    def wrapper(*args, **kwargs):
        start = time.time()
        result = fun(*args, **kwargs)
        end = time.time()
        run = end-start
        print(f'运行时长：{run}')
        return result
    return wrapper


@runtime
def runfunc():
    time.sleep(1)


runfunc()	# 运行时长：1.0027356147766113

2. 使用装饰器记录日志

import functools
import logging

# 日志配置
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')


def my_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        logging.info(f'{func.__name__} 函数输入参数 args：{args}，kwargs：{kwargs}')
        result = func(*args, **kwargs)
        if isinstance(result, (int, float)):
            logging.info(f'{func.__name__} 函数执行结果：{result}')
        elif isinstance(result, str):
            logging.error(f'{func.__name__} 函数执行结果：{result}')
        return result
    return wrapper


@my_decorator
def example(x):
    try:
        result = 100 / x
    except ZeroDivisionError as e:
        return f'发生错误 "{e}"'
    except TypeError as e:
        return f'发生错误 "{e}"'
    return result


example(1)
example(1.2)
example(0)
example('a')

输出结果

2024-09-09 17:01:35,856 - INFO - example 函数输入参数 args：(1,)，kwargs：{}
2024-09-09 17:01:35,856 - INFO - example 函数执行结果：100.0
2024-09-09 17:01:35,856 - INFO - example 函数输入参数 args：(1.2,)，kwargs：{}
2024-09-09 17:01:35,856 - INFO - example 函数执行结果：83.33333333333334
2024-09-09 17:01:35,856 - INFO - example 函数输入参数 args：(0,)，kwargs：{}
2024-09-09 17:01:35,856 - ERROR - example 函数执行结果：发生错误 "division by zero"
2024-09-09 17:01:35,856 - INFO - example 函数输入参数 args：('a',)，kwargs：{}
2024-09-09 17:01:35,856 - ERROR - example 函数执行结果：发生错误"unsupported operand type(s) for /: 'int' and 'str'"