Python 装饰器高级应用指南

Python 装饰器高级应用指南

1. 什么是装饰器？

装饰器是 Python 中一种特殊的语法结构，用于修改函数或类的行为。它允许我们在不修改原函数代码的情况下，为函数添加额外的功能。

2. 基本语法

装饰器使用 @ 符号来应用，放在函数定义的上方。

@decorator
def function():
    pass

3. 装饰器的实现

3.1 简单装饰器

def my_decorator(func):
    def wrapper():
        print("Before function call")
        func()
        print("After function call")
    return wrapper

@my_decorator
def say_hello():
    print("Hello")

say_hello()
# 输出:
# Before function call
# Hello
# After function call

3.2 带参数的装饰器

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

@repeat(3)
def say_hello(name):
    print(f"Hello, {name}!")

say_hello("Alice")
# 输出:
# Hello, Alice!
# Hello, Alice!
# Hello, Alice!

3.3 保留函数元数据

使用 functools.wraps 可以保留原函数的元数据，如函数名、文档字符串等。

import functools

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

@my_decorator
def say_hello():
    """Say hello function"""
    print("Hello")

print(say_hello.__name__)  # 输出: say_hello
print(say_hello.__doc__)   # 输出: Say hello function

4. 高级应用

4.1 类装饰器

类装饰器是使用类来实现的装饰器，它通过 __call__ 方法来实现装饰功能。

class CountCalls:
    def __init__(self, func):
        self.func = func
        self.count = 0
    
    def __call__(self, *args, **kwargs):
        self.count += 1
        print(f"Call {self.count} of {self.func.__name__}")
        return self.func(*args, **kwargs)

@CountCalls
def say_hello():
    print("Hello")

say_hello()  # 输出: Call 1 of say_hello
             #       Hello
say_hello()  # 输出: Call 2 of say_hello
             #       Hello

4.2 装饰器链

多个装饰器可以同时应用到一个函数上，形成装饰器链。

def decorator1(func):
    def wrapper(*args, **kwargs):
        print("Decorator 1 before")
        result = func(*args, **kwargs)
        print("Decorator 1 after")
        return result
    return wrapper

def decorator2(func):
    def wrapper(*args, **kwargs):
        print("Decorator 2 before")
        result = func(*args, **kwargs)
        print("Decorator 2 after")
        return result
    return wrapper

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

say_hello()
# 输出:
# Decorator 1 before
# Decorator 2 before
# Hello
# Decorator 2 after
# Decorator 1 after

4.3 带状态的装饰器

装饰器可以保持状态，例如缓存函数结果。

import functools

def cache(func):
    """Cache function results"""
    cache_dict = {}
    
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        key = str(args) + str(kwargs)
        if key not in cache_dict:
            cache_dict[key] = func(*args, **kwargs)
        return cache_dict[key]
    
    return wrapper

@cache
def fibonacci(n):
    if n <= 1:
        return n
    return fibonacci(n-1) + fibonacci(n-2)

print(fibonacci(30))  # 第一次计算，会缓存结果
print(fibonacci(30))  # 直接从缓存获取结果

4.4 装饰器工厂

装饰器工厂是返回装饰器的函数，允许我们为装饰器提供配置参数。

def log_level(level):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            print(f"[{level.upper()}] Calling {func.__name__}")
            result = func(*args, **kwargs)
            print(f"[{level.upper()}] {func.__name__} completed")
            return result
        return wrapper
    return decorator

@log_level("info")
def process_data(data):
    print(f"Processing data: {data}")
    return data * 2

@log_level("error")
def handle_error():
    print("Handling error")

process_data(42)  # 输出: [INFO] Calling process_data
                  #       Processing data: 42
                  #       [INFO] process_data completed

handle_error()    # 输出: [ERROR] Calling handle_error
                  #       Handling error
                  #       [ERROR] handle_error completed

5. 实际应用场景

5.1 日志记录

def log_function(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print(f"Calling {func.__name__} with args: {args}, kwargs: {kwargs}")
        result = func(*args, **kwargs)
        print(f"{func.__name__} returned: {result}")
        return result
    return wrapper

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

add(3, 5)  # 输出调用信息和返回值

5.2 性能计时

import time

def timer(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        print(f"{func.__name__} took {end_time - start_time:.4f} seconds")
        return result
    return wrapper

@timer
def slow_function():
    time.sleep(1)
    return "Done"

slow_function()  # 输出执行时间

5.3 权限验证

def require_permission(permission):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            # 假设当前用户权限存储在 current_user.permissions 中
            current_user = {"permissions": ["read", "write"]}
            if permission not in current_user["permissions"]:
                raise PermissionError(f"Required permission: {permission}")
            return func(*args, **kwargs)
        return wrapper
    return decorator

@require_permission("write")
def create_post(content):
    print(f"Creating post: {content}")
    return "Post created"

@require_permission("admin")
def delete_post(post_id):
    print(f"Deleting post: {post_id}")
    return "Post deleted"

create_post("Hello World")  # 成功执行
try:
    delete_post(1)
except PermissionError as e:
    print(e)  # 抛出权限错误

5.4 输入验证

def validate_input(*types):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            for arg, type_ in zip(args, types):
                if not isinstance(arg, type_):
                    raise TypeError(f"Expected {type_}, got {type(arg)}")
            return func(*args, **kwargs)
        return wrapper
    return decorator

@validate_input(int, str)
def process_user(user_id, name):
    print(f"Processing user {user_id}: {name}")
    return f"User {user_id} processed"

process_user(1, "Alice")  # 成功执行
try:
    process_user("1", "Alice")
except TypeError as e:
    print(e)  # 抛出类型错误

6. 最佳实践

1. 使用 functools.wraps：保留原函数的元数据。
2. 保持装饰器简单：每个装饰器只负责一项功能。
3. 使用装饰器工厂：当需要为装饰器提供参数时。
4. 文档化装饰器：为装饰器添加清晰的文档字符串。
5. 测试装饰器：确保装饰器在各种情况下都能正常工作。

7. 总结

装饰器是 Python 中非常强大的功能，它允许我们以一种简洁、优雅的方式修改函数或类的行为。通过掌握装饰器的高级应用，我们可以编写更加模块化、可维护的代码。

在实际应用中，装饰器可以用于日志记录、性能计时、权限验证、输入验证等多种场景，大大提高了代码的可读性和可维护性。

希望本文对你理解和应用 Python 装饰器有所帮助！