python使用装饰器进行函数增强

2023-03-30 00:00:00 函数 装饰 增强

在 Python 中，装饰器（Decorator）是一种语法，允许在不改变函数本身的情况下增强或修改函数的行为。装饰器本质上是一个函数，可以接收一个函数作为参数，并返回一个新的函数。

下面是一个简单的装饰器示例，该装饰器可以在函数执行前后打印一些信息：

def my_decorator(func):
    def wrapper(*args, **kwargs):
        print("Before the function is called.")
        result = func(*args, **kwargs)
        print("After the function is called.")
        return result
    return wrapper

@my_decorator
def my_function():
    print("Inside the function.")

my_function()

在上面的示例中，my_decorator 是一个装饰器函数，它接受一个函数 func 作为参数，并返回一个新的函数 wrapper。wrapper 函数接受任意数量的位置参数和关键字参数，并在执行被装饰函数之前和之后打印一些信息。

@my_decorator 表示将装饰器应用于 my_function 函数。因此，当调用 my_function() 时，实际上是调用了装饰器函数 my_decorator 返回的 wrapper 函数。wrapper 函数在执行被装饰函数 my_function 前后会打印信息，并返回被装饰函数的结果。

下面是一个更具体的例子，它演示了如何使用装饰器来检查函数参数的类型是否正确：

def check_types(*arg_types, **kwarg_types):
    def decorator(func):
        def wrapper(*args, **kwargs):
            for arg, arg_type in zip(args, arg_types):
                if not isinstance(arg, arg_type):
                    raise TypeError(f"Expected {arg_type} but got {type(arg)}")
            for arg_name, arg_type in kwarg_types.items():
                if arg_name in kwargs and not isinstance(kwargs[arg_name], arg_type):
                    raise TypeError(f"Expected {arg_type} for {arg_name} but got {type(kwargs[arg_name])}")
            return func(*args, **kwargs)
        return wrapper
    return decorator

@check_types(int, str, str)
def say_hello(age, name, company):
    print(f"Hello, my name is {name} and I work at {company}. I am {age} years old.")

say_hello(30, "pidancode.com", "皮蛋编程")

在上面的示例中，check_types 是一个装饰器工厂函数，它接受任意数量的参数类型和关键字参数类型，并返回一个装饰器函数 decorator。decorator 函数接受一个函数 func 作为参数，并返回一个新的函数 wrapper。wrapper 函数首先检查所有位置参数和关键字参数的类型是否正确，如果类型不正确则抛出一个 TypeError 异常，否则执行被装饰函数并返回结果。

@check_types(int, str, str) 表示将装饰器应用于 say_hello函数。装饰器工厂函数 check_types 的参数类型为 int, str, str，表示 say_hello 函数的前三个参数分别为整型、字符串和字符串。

当调用 say_hello(30, "pidancode.com", "皮蛋编程") 时，实际上是调用了装饰器函数 check_types(int, str, str) 返回的装饰器函数 decorator，decorator 函数又返回一个新的函数 wrapper，wrapper 函数首先检查参数类型是否正确，然后执行被装饰函数 say_hello 并返回结果。

注意，在使用装饰器时，被装饰的函数的元信息（如函数名、文档字符串等）会被装饰器函数所取代。为了解决这个问题，可以使用 functools.wraps 装饰器来将被装饰函数的元信息复制到装饰器函数中。例如：

import functools

def my_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print("Before the function is called.")
        result = func(*args, **kwargs)
        print("After the function is called.")
        return result
    return wrapper

@my_decorator
def my_function():
    """This is a docstring."""
    print("Inside the function.")

print(my_function.__name__)  # 输出 "my_function"
print(my_function.__doc__)   # 输出 "This is a docstring."

在上面的示例中，使用 functools.wraps 装饰器将被装饰函数的元信息复制到装饰器函数中。my_function 函数的 name 属性和 doc 属性分别为 "my_function" 和 "This is a docstring."。

相关文章