python之装饰器

装饰器是 Python 中的一种高级特性，用于在不修改函数或方法定义的情况下，动态地增加功能。装饰器本质上是一个函数，它接受一个函数作为参数，并返回一个新的函数。装饰器通常用于日志记录、性能测试、事务处理、权限校验等场景。

1. 基本装饰器

定义和使用装饰器

def my_decorator(func):
    def wrapper():
        print("Something is happening before the function is called.")
        func()
        print("Something is happening after the function is called.")
    return wrapper

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

# 调用被装饰的函数
say_hello()
# 输出：
# Something is happening before the function is called.
# Hello!
# Something is happening after the function is called.

2. 带参数的装饰器

处理带参数的函数

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
    return wrapper

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

# 调用被装饰的函数
say_hello("Alice")
# 输出：
# Something is happening before the function is called.
# Hello, Alice!
# Something is happening after the function is called.

3. 带参数的装饰器

有时我们需要给装饰器本身传递参数，这时可以使用一个外层函数来接受参数，并返回一个装饰器。

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 say_hello(name):
    print(f"Hello, {name}!")

# 调用被装饰的函数
say_hello("Alice")
# 输出：
# Hello, Alice!
# Hello, Alice!
# Hello, Alice!

4. 类装饰器

装饰器也可以用类来实现。类装饰器需要实现 __call__ 方法。

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

    def __call__(self, *args, **kwargs):
        print("Something is happening before the function is called.")
        result = self.func(*args, **kwargs)
        print("Something is happening after the function is called.")
        return result

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

# 调用被装饰的函数
say_hello("Alice")
# 输出：
# Something is happening before the function is called.
# Hello, Alice!
# Something is happening after the function is called.

5. 多个装饰器

一个函数可以被多个装饰器装饰，装饰器的应用顺序是从内到外。

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

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

@decorator1
@decorator2
def say_hello(name):
    print(f"Hello, {name}!")

# 调用被装饰的函数
say_hello("Alice")
# 输出：
# Decorator 1
# Decorator 2
# Hello, Alice!

6. 保留原函数的元数据

使用 functools.wraps 来保留原函数的元数据时，装饰器不会改变原函数的名称、文档字符串等属性。下面是你提供的代码以及调用 say_hello("Alice") 时的输出：

import functools

def my_decorator(func):
    @functools.wraps(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
    return wrapper

@my_decorator
def say_hello(name):
    """This is a greeting function."""
    print(f"Hello, {name}!")

# 调用被装饰的函数
say_hello("Alice")

调用 say_hello("Alice") 时的输出将是：

Something is happening before the function is called.
Hello, Alice!
Something is happening after the function is called.

此外，由于使用了 functools.wraps，原函数的元数据（如函数名和文档字符串）得以保留。你可以通过以下代码验证这一点：

print(say_hello.__name__)  # 输出: say_hello
print(say_hello.__doc__)   # 输出: This is a greeting function.

如果没有使用 functools.wraps，则 say_hello.__name__ 会变成 wrapper，say_hello.__doc__ 会变成 None 或者是 wrapper 函数的文档字符串。使用 functools.wraps 可以确保这些元数据保持不变。

使用场景

1. 日志记录

日志记录是开发和维护应用程序的重要部分。通过装饰器，可以在函数调用前后记录日志信息。

import logging

# 配置日志
logging.basicConfig(level=logging.INFO)

def log_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        logging.info(f"Calling function {func.__name__} with args: {args} and kwargs: {kwargs}")
        result = func(*args, **kwargs)
        logging.info(f"Function {func.__name__} returned {result}")
        return result
    return wrapper

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

# 调用被装饰的函数
result = add(3, 5)
# 输出：
# INFO:root:Calling function add with args: (3, 5) and kwargs: {}
# INFO:root:Function add returned 8

2. 性能测试

性能测试可以帮助我们了解函数的执行时间，从而进行优化。通过装饰器，可以在函数执行前后记录时间差。

import time
import functools

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

@timing_decorator
def slow_function():
    time.sleep(2)
    return "Finished"

# 调用被装饰的函数
result = slow_function()
# 输出：
# Function slow_function took 2.0001 seconds to execute

3. 事务处理

在数据库操作中，事务处理是确保数据一致性的重要机制。通过装饰器，可以在函数执行前后管理事务。

import functools

class Database:
    def __init__(self):
        self.transaction_active = False

    def start_transaction(self):
        self.transaction_active = True
        print("Transaction started")

    def commit_transaction(self):
        self.transaction_active = False
        print("Transaction committed")

    def rollback_transaction(self):
        self.transaction_active = False
        print("Transaction rolled back")

db = Database()

def transaction_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        db.start_transaction()
        try:
            result = func(*args, **kwargs)
            db.commit_transaction()
            return result
        except Exception as e:
            db.rollback_transaction()
            raise e
    return wrapper

@transaction_decorator
def perform_db_operations():
    print("Performing database operations")
    # 模拟数据库操作
    if True:  # 模拟操作成功
        return "Success"
    else:
        raise Exception("Database error")

# 调用被装饰的函数
result = perform_db_operations()
# 输出：
# Transaction started
# Performing database operations
# Transaction committed

4. 权限校验

在 Web 应用中，权限校验是确保用户只能访问其有权限访问的资源的重要机制。通过装饰器，可以在函数执行前进行权限校验。

import functools

def check_permission(user_role):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            if user_role != "admin":
                raise PermissionError("You do not have permission to access this resource")
            return func(*args, **kwargs)
        return wrapper
    return decorator

@check_permission(user_role="admin")
def delete_user(user_id):
    print(f"User {user_id} deleted")

# 调用被装饰的函数
try:
    delete_user(123)
    # 输出：
    # User 123 deleted
except PermissionError as e:
    print(e)

@check_permission(user_role="guest")
def delete_user(user_id):
    print(f"User {user_id} deleted")

# 调用被装饰的函数
try:
    delete_user(123)
    # 输出：
    # You do not have permission to access this resource
except PermissionError as e:
    print(e)