MK
摩柯社区 - 一个极简的技术知识社区
AI 面试

Python装饰器的适用场景

2021-10-095.3k 阅读

日志记录

在软件开发过程中，日志记录是一项至关重要的任务。它有助于开发人员在程序运行时跟踪和诊断问题。Python装饰器为实现日志记录提供了一种简洁且高效的方式。

函数调用日志

假设我们有一个简单的函数，用于计算两个数的和：

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

现在，我们希望每次调用这个函数时，都能记录下函数的参数和返回值。我们可以使用装饰器来实现这一点：

import logging

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

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

在上述代码中，log_function_call 是一个装饰器。它接受一个函数 func 作为参数，并返回一个新的函数 wrapper。wrapper 函数在调用原始函数 func 前后，使用 logging 模块记录函数调用的相关信息。

当我们调用 add(2, 3) 时，日志中会记录如下信息：

INFO:root:Calling function add with args: (2, 3), kwargs: {}
INFO:root:Function add returned: 5

这种方式不仅清晰地记录了函数调用的细节，而且不会污染原始函数的代码逻辑。原始的 add 函数依然保持简洁，只专注于其核心功能——加法运算。

异常日志记录

除了记录函数的正常调用，我们还可以使用装饰器记录函数执行过程中发生的异常。例如，假设我们有一个函数，可能会引发 ZeroDivisionError：

def divide(a, b):
    return a / b

我们可以创建一个装饰器来捕获并记录异常：

import logging

def log_exceptions(func):
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            logging.error(f"An error occurred in {func.__name__}: {e}")
    return wrapper

@log_exceptions
def divide(a, b):
    return a / b

当我们调用 divide(10, 0) 时，日志中会记录如下信息：

ERROR:root:An error occurred in divide: division by zero

这样，开发人员可以通过查看日志，快速定位程序运行过程中出现的异常，大大提高了调试效率。

性能分析

在优化代码性能时，了解函数的执行时间是非常重要的。Python装饰器可以方便地实现对函数执行时间的测量。

测量单个函数执行时间

假设我们有一个计算斐波那契数列的函数：

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

为了测量这个函数的执行时间，我们可以创建一个装饰器：

import time

def measure_execution_time(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        execution_time = end_time - start_time
        print(f"Function {func.__name__} took {execution_time} seconds to execute.")
        return result
    return wrapper

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

当我们调用 fibonacci(30) 时，会输出函数的执行时间：

Function fibonacci took 12.586234724044799 seconds to execute.

通过这种方式，我们可以直观地了解到函数的运行效率，以便对其进行针对性的优化。

比较不同实现的性能

有时候，我们可能有多种方法来解决同一个问题，需要比较不同方法的性能。例如，对于计算斐波那契数列，我们还可以使用动态规划的方法来优化：

def fibonacci_dp(n):
    if n <= 1:
        return n
    dp = [0] * (n + 1)
    dp[1] = 1
    for i in range(2, n + 1):
        dp[i] = dp[i - 1] + dp[i - 2]
    return dp[n]

我们可以使用装饰器分别测量这两个函数的执行时间：

import time

def measure_execution_time(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        execution_time = end_time - start_time
        print(f"Function {func.__name__} took {execution_time} seconds to execute.")
        return result
    return wrapper

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

@measure_execution_time
def fibonacci_dp(n):
    if n <= 1:
        return n
    dp = [0] * (n + 1)
    dp[1] = 1
    for i in range(2, n + 1):
        dp[i] = dp[i - 1] + dp[i - 2]
    return dp[n]

当我们分别调用 fibonacci(30) 和 fibonacci_dp(30) 时，可以看到明显的性能差异：

Function fibonacci took 12.586234724044799 seconds to execute.
Function fibonacci_dp took 0.0000019073486328125 seconds to execute.

这清楚地表明了动态规划实现的斐波那契函数在性能上远远优于递归实现。

权限控制

在开发Web应用或其他需要权限管理的系统时，确保只有授权的用户能够访问特定的功能是至关重要的。Python装饰器可以有效地实现权限控制。

简单的用户权限检查

假设我们有一个用户类和一些需要权限控制的函数：

class User:
    def __init__(self, username, role):
        self.username = username
        self.role = role

def is_admin(user):
    return user.role == 'admin'

def restricted_access(func):
    def wrapper(user, *args, **kwargs):
        if is_admin(user):
            return func(user, *args, **kwargs)
        else:
            print("Access denied. You do not have sufficient permissions.")
    return wrapper

@restricted_access
def delete_user(user, target_user):
    print(f"{user.username} is deleting user {target_user.username}")

在上述代码中，restricted_access 装饰器检查用户是否具有管理员权限。如果用户是管理员，则允许执行被装饰的函数 delete_user；否则，提示访问被拒绝。

我们可以这样调用：

admin_user = User('admin', 'admin')
regular_user = User('user', 'user')

delete_user(admin_user, regular_user)
delete_user(regular_user, admin_user)

输出结果为：

admin is deleting user user
Access denied. You do not have sufficient permissions.

这样，通过装饰器，我们可以轻松地对不同权限的用户进行功能访问控制。

基于角色的权限管理

在更复杂的系统中，可能存在多种角色，不同角色具有不同的权限。我们可以扩展权限检查的逻辑。假设我们有三种角色：管理员（admin）、经理（manager）和普通用户（user），并且管理员可以执行所有操作，经理可以执行部分操作，普通用户只能执行有限操作。

class User:
    def __init__(self, username, role):
        self.username = username
        self.role = role

def has_permission(user, action):
    if user.role == 'admin':
        return True
    elif user.role =='manager':
        if action in ['view_report', 'approve_request']:
            return True
        return False
    elif user.role == 'user':
        if action in ['view_profile']:
            return True
        return False
    return False

def role_based_permission(action):
    def decorator(func):
        def wrapper(user, *args, **kwargs):
            if has_permission(user, action):
                return func(user, *args, **kwargs)
            else:
                print("Access denied. You do not have sufficient permissions.")
        return wrapper
    return decorator

@role_based_permission('delete_user')
def delete_user(user, target_user):
    print(f"{user.username} is deleting user {target_user.username}")

@role_based_permission('view_report')
def view_report(user):
    print(f"{user.username} is viewing the report")

admin_user = User('admin', 'admin')
manager_user = User('manager','manager')
regular_user = User('user', 'user')

delete_user(admin_user, regular_user)
delete_user(manager_user, regular_user)
view_report(manager_user)
view_report(regular_user)

输出结果为：

admin is deleting user user
Access denied. You do not have sufficient permissions.
manager is viewing the report
Access denied. You do not have sufficient permissions.

这里，role_based_permission 装饰器工厂函数接受一个操作名称作为参数，然后返回一个装饰器。这个装饰器根据用户的角色和操作名称来检查权限，实现了更灵活的基于角色的权限管理。

输入验证

在函数调用之前，对输入参数进行验证是确保程序健壮性的重要步骤。Python装饰器可以方便地实现输入验证逻辑。

基本类型检查

假设我们有一个函数，用于计算两个整数的乘积：

def multiply(a, b):
    return a * b

我们希望确保输入的参数都是整数类型。可以使用装饰器来实现类型检查：

def validate_type(expected_type):
    def decorator(func):
        def wrapper(*args, **kwargs):
            for arg in args:
                if not isinstance(arg, expected_type):
                    raise TypeError(f"Expected {expected_type.__name__}, got {type(arg).__name__}")
            for value in kwargs.values():
                if not isinstance(value, expected_type):
                    raise TypeError(f"Expected {expected_type.__name__}, got {type(value).__name__}")
            return func(*args, **kwargs)
        return wrapper
    return decorator

@validate_type(int)
def multiply(a, b):
    return a * b

当我们调用 multiply(2, 3) 时，一切正常。但如果调用 multiply(2, '3')，会抛出 TypeError：

TypeError: Expected int, got str

通过这种方式，我们可以在函数执行前捕获类型错误，避免在函数内部出现难以调试的问题。

复杂输入验证

对于更复杂的输入验证，比如验证一个字符串是否符合特定的格式（例如邮箱格式）。假设我们有一个函数，用于向指定邮箱发送消息：

import re

def send_email(email, message):
    print(f"Sending message '{message}' to {email}")

def validate_email(func):
    def wrapper(email, *args, **kwargs):
        if not re.match(r"[^@]+@[^@]+\.[^@]+", email):
            raise ValueError("Invalid email address")
        return func(email, *args, **kwargs)
    return wrapper

@validate_email
def send_email(email, message):
    print(f"Sending message '{message}' to {email}")

当我们调用 send_email('valid@example.com', 'Hello') 时，会正常输出：

Sending message 'Hello' to valid@example.com

但如果调用 send_email('invalidemail', 'Hello')，会抛出 ValueError：

ValueError: Invalid email address

这样，通过装饰器实现的复杂输入验证，能够有效地保证函数在正确的输入下运行，提高程序的稳定性和可靠性。

缓存

在一些计算密集型的应用中，重复计算相同的结果会浪费大量的时间和资源。Python装饰器可以实现缓存机制，避免重复计算。

简单函数结果缓存

假设我们有一个计算阶乘的函数：

def factorial(n):
    if n == 0 or n == 1:
        return 1
    else:
        return n * factorial(n - 1)

为了避免重复计算相同参数的阶乘，我们可以使用装饰器来实现缓存：

def cache_result(func):
    cache = {}
    def wrapper(*args, **kwargs):
        key = (args, tuple(sorted(kwargs.items())))
        if key in cache:
            return cache[key]
        result = func(*args, **kwargs)
        cache[key] = result
        return result
    return wrapper

@cache_result
def factorial(n):
    if n == 0 or n == 1:
        return 1
    else:
        return n * factorial(n - 1)

在上述代码中，cache_result 装饰器使用一个字典 cache 来存储已经计算过的函数结果。每次调用函数时，先检查参数对应的结果是否已经在缓存中，如果存在则直接返回缓存中的结果，否则计算结果并缓存起来。

这样，当我们多次调用 factorial(5) 时，实际的计算过程只会执行一次，后续调用直接从缓存中获取结果，大大提高了函数的执行效率。

缓存过期机制

在某些情况下，我们可能希望缓存的结果在一段时间后过期，需要重新计算。例如，假设我们有一个函数，用于获取实时股票价格：

import time

def get_stock_price(stock_symbol):
    # 模拟获取实时股票价格的逻辑
    print(f"Fetching real - time price for {stock_symbol}")
    return 100  # 假设返回价格为100

def cache_with_expiry(expiry_time):
    def decorator(func):
        cache = {}
        timestamps = {}
        def wrapper(*args, **kwargs):
            key = (args, tuple(sorted(kwargs.items())))
            current_time = time.time()
            if key in cache and current_time - timestamps[key] < expiry_time:
                return cache[key]
            result = func(*args, **kwargs)
            cache[key] = result
            timestamps[key] = current_time
            return result
        return wrapper
    return decorator

@cache_with_expiry(60)  # 缓存60秒
def get_stock_price(stock_symbol):
    # 模拟获取实时股票价格的逻辑
    print(f"Fetching real - time price for {stock_symbol}")
    return 100  # 假设返回价格为100

在上述代码中，cache_with_expiry 装饰器不仅使用 cache 字典存储结果，还使用 timestamps 字典记录每个结果的缓存时间。每次调用函数时，检查缓存是否过期，如果未过期则返回缓存结果，否则重新计算并更新缓存和时间戳。

这样，我们既利用了缓存提高了性能，又保证了缓存结果不会长期使用而过时，适用于需要获取实时数据但又希望减少频繁获取开销的场景。

上下文管理

上下文管理在Python中常用于资源的分配和释放，比如文件的打开和关闭、数据库连接的管理等。Python装饰器可以以一种简洁的方式实现上下文管理的功能。

模拟文件操作的上下文管理

假设我们有一个函数，用于读取文件内容：

def read_file(file_path):
    with open(file_path, 'r') as file:
        return file.read()

我们可以使用装饰器来实现类似的上下文管理逻辑，而不使用 with 语句：

def file_context_manager(func):
    def wrapper(file_path, *args, **kwargs):
        file = open(file_path, 'r')
        try:
            result = func(file, *args, **kwargs)
            return result
        finally:
            file.close()
    return wrapper

@file_context_manager
def read_file(file):
    return file.read()

在上述代码中，file_context_manager 装饰器在调用被装饰的函数 read_file 前打开文件，在函数执行完毕后关闭文件。这样，即使在函数执行过程中发生异常，文件也能被正确关闭，避免了资源泄漏。

数据库连接的上下文管理

在数据库操作中，正确管理数据库连接同样重要。假设我们使用 sqlite3 模块进行数据库操作：

import sqlite3

def execute_query(query):
    conn = sqlite3.connect('example.db')
    cursor = conn.cursor()
    cursor.execute(query)
    result = cursor.fetchall()
    conn.close()
    return result

我们可以使用装饰器来简化数据库连接的管理：

def db_connection_manager(func):
    def wrapper(*args, **kwargs):
        conn = sqlite3.connect('example.db')
        try:
            cursor = conn.cursor()
            result = func(cursor, *args, **kwargs)
            conn.commit()
            return result
        except Exception as e:
            conn.rollback()
            raise e
        finally:
            conn.close()
    return wrapper

@db_connection_manager
def execute_query(cursor, query):
    cursor.execute(query)
    return cursor.fetchall()

在这个例子中，db_connection_manager 装饰器负责建立数据库连接、创建游标、执行被装饰的函数（这里是执行SQL查询）、提交事务（如果执行成功）、回滚事务（如果发生异常）以及关闭数据库连接。通过这种方式，将数据库连接的管理逻辑从业务逻辑（执行查询）中分离出来，使得代码更加清晰和易于维护。

函数注册

在一些框架或库的开发中，需要将函数注册到某个全局的注册表中，以便后续统一调用或管理。Python装饰器可以方便地实现函数注册功能。

简单的函数注册

假设我们有一个模块，其中定义了一些处理不同类型任务的函数，我们希望将这些函数注册到一个任务注册表中。

task_registry = {}

def register_task(func):
    task_registry[func.__name__] = func
    return func

@register_task
def task1():
    print("Executing task1")

@register_task
def task2():
    print("Executing task2")

在上述代码中，register_task 装饰器将被装饰的函数添加到 task_registry 字典中，键为函数名，值为函数对象本身。这样，我们可以通过访问 task_registry 来获取和调用注册的函数。例如：

if 'task1' in task_registry:
    task_registry['task1']()

输出结果为：

Executing task1

通过这种方式，我们可以轻松地管理和调用不同的任务函数，适用于任务调度、插件系统等场景。

带参数的函数注册

有时候，我们可能需要在注册函数时传递一些额外的参数，以便对注册的函数进行更细粒度的控制。例如，假设我们有一个任务处理系统，每个任务可能有不同的优先级。

task_registry = {}

def register_task(priority):
    def decorator(func):
        task_registry[func.__name__] = {'function': func, 'priority': priority}
        return func
    return decorator

@register_task(priority = 1)
def task1():
    print("Executing task1")

@register_task(priority = 2)
def task2():
    print("Executing task2")

在这个例子中，register_task 是一个装饰器工厂函数，它接受一个 priority 参数。返回的装饰器将函数及其优先级信息注册到 task_registry 中。我们可以根据优先级对注册的任务进行排序和调度。例如：

sorted_tasks = sorted(task_registry.items(), key = lambda item: item[1]['priority'])
for task_name, task_info in sorted_tasks:
    task_info['function']()

输出结果为：

Executing task1
Executing task2

这样，通过带参数的函数注册，我们可以实现更灵活的函数管理和调度机制。

类方法装饰

在Python类中，装饰器同样有着广泛的应用。它可以用于修改类方法的行为，实现诸如延迟加载、代理等功能。

延迟加载属性

假设我们有一个类，其中某个属性的计算比较耗时，我们希望在实际访问该属性时才进行计算，而不是在类实例化时就计算。

class LazyLoadClass:
    def __init__(self):
        self._expensive_attribute = None

    def _calculate_expensive_attribute(self):
        # 模拟耗时计算
        import time
        time.sleep(2)
        return 42

    @property
    def expensive_attribute(self):
        if self._expensive_attribute is None:
            self._expensive_attribute = self._calculate_expensive_attribute()
        return self._expensive_attribute

在上述代码中，我们使用 property 装饰器将 expensive_attribute 方法转换为一个属性。当我们访问 instance.expensive_attribute 时，会首先检查 _expensive_attribute 是否已经计算，如果没有则调用 _calculate_expensive_attribute 方法进行计算。这样，实现了属性的延迟加载，提高了类实例化的速度。

方法代理

有时候，我们希望在类的某个方法调用前或调用后执行一些额外的逻辑，而不改变方法的核心功能。例如，假设我们有一个类，其中的方法需要在调用前后记录日志。

import logging

class MethodProxyClass:
    def __init__(self):
        pass

    def log_method_call(func):
        def wrapper(self, *args, **kwargs):
            logging.info(f"Calling method {func.__name__}")
            result = func(self, *args, **kwargs)
            logging.info(f"Method {func.__name__} called")
            return result
        return wrapper

    @log_method_call
    def example_method(self):
        print("Executing example method")

在上述代码中，log_method_call 装饰器在调用 example_method 前后记录日志。这种方法代理的方式可以在不修改 example_method 核心代码的情况下，为其添加额外的功能，提高了代码的可维护性和可扩展性。

通过以上各种场景的介绍，我们可以看到Python装饰器在代码开发中具有强大的功能和广泛的应用。它能够有效地分离关注点，提高代码的复用性、可维护性和可读性，是Python开发者不可或缺的工具之一。无论是在小型脚本还是大型项目中，合理使用装饰器都能为代码带来显著的优化和提升。