嵌套函数中巧用装饰器避免重复打印

问题描述

在Python开发中，装饰器（Decorator）是实现横切关注点（如日志、性能监控、权限验证等）的强大工具。然而，当多个函数都应用了同一个装饰器，并且这些函数之间存在嵌套调用关系时，可能会出现意料之外的重复输出。

例如，我们有一个简单的计时装饰器 @time_elapsed，用于测量函数的执行时间并打印结果。

import time
from functools import wraps

def time_elapsed(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        elapsed_time = time.time() - start_time
        print(f'{func.__name__} took {elapsed_time:.2f} seconds.')
        return result
    return wrapper

@time_elapsed
def func1():
    time.sleep(0.1)

@time_elapsed
def func2():
    func1()
    time.sleep(0.2)

当我们独立调用 func1() 时，输出符合预期：

func1 took 0.10 seconds.

然而，当我们调用 func2() 时，由于 func2 内部调用了 func1，并且 func1 也被 @time_elapsed 装饰，导致 func1 的计时信息被打印了两次（一次作为独立调用，一次作为 func2 的子调用），这通常不是我们希望看到的：

func1 took 0.10 seconds.  # func2 内部调用 func1 产生的输出
func2 took 0.30 seconds.

我们的目标是，当调用 func2() 时，只打印 func2 的计时信息，即：

func2 took 0.30 seconds.

同时，func1() 独立调用时仍能正常打印其计时信息。

解决方案：基于深度计数的装饰器控制

为了解决上述问题，我们可以在装饰器内部引入一个机制来跟踪当前函数调用的深度。通过维护一个全局或装饰器级别的计数器，我们可以判断当前执行的函数是否是最外层的被装饰函数调用，或者是否达到了我们希望打印输出的特定深度。

修改后的 time_elapsed 装饰器将包含一个内部计数器 _timer_running 和一个深度阈值 DEPTH。

import time
from functools import wraps

def time_elapsed(func):
    # 定义打印输出的深度。DEPTH=1 表示只打印最外层调用。
    # 可以通过修改此值来控制打印的嵌套层级。
    DEPTH = 1 

    # 初始化一个装饰器级别的计数器，用于跟踪当前函数调用的嵌套深度。
    # 首次调用时，time_elapsed._timer_running 不存在，设置为0。
    if not hasattr(time_elapsed, '_timer_running'):
        time_elapsed._timer_running = 0

    @wraps(func)
    def wrapper(*args, **kwargs):
        # 如果当前嵌套深度大于等于设定的DEPTH，则跳过计时和打印。
        # 这意味着我们只关心特定深度内的函数计时。
        if time_elapsed._timer_running >= DEPTH:
            return func(*args, **kwargs)

        # 如果当前深度小于DEPTH，则需要进行计时。
        # 在执行函数前，增加计数器，表示进入了一个新的计时层级。
        time_elapsed._timer_running += 1

        try:
            # 执行原始函数并计时
            start_time = time.time()
            result = func(*args, **kwargs)
            elapsed_time = time.time() - start_time
            print(f'{func.__name__} took {elapsed_time:.2f} seconds.')
        finally:
            # 无论函数执行成功与否，在函数退出时，都需减少计数器。
            # 确保计数器正确回溯，避免影响后续的独立调用。
            time_elapsed._timer_running -= 1

        return result

    return wrapper

# 示例函数保持不变
@time_elapsed
def func1():
    time.sleep(0.1)

@time_elapsed
def func2():
    func1()
    time.sleep(0.2)

@time_elapsed
def func3():
    func1()
    func2()
    time.sleep(0.3)

@time_elapsed
def func4():
    func1()
    func2()
    func3()
    time.sleep(0.4)

if __name__ == "__main__":
    print("--- Testing func1 ---")
    func1()
    print("\n--- Testing func2 ---")
    func2()
    print("\n--- Testing func3 ---")
    func3()
    print("\n--- Testing func4 ---")
    func4()

运行效果与解释

当 DEPTH = 1 时，运行上述代码，我们将得到以下输出：

--- Testing func1 ---
func1 took 0.10 seconds.

--- Testing func2 ---
func2 took 0.30 seconds.

--- Testing func3 ---
func3 took 0.70 seconds.

--- Testing func4 ---
func4 took 1.50 seconds.

解释：

1. func1() 调用: _timer_running 为 0。小于 DEPTH (1)，因此 _timer_running 增至 1，执行计时和打印，然后减至 0。
2. func2() 调用:
   • 外部 func2 调用：_timer_running 为 0。小于 DEPTH (1)，_timer_running 增至 1。
   • 内部 func1 调用：此时 _timer_running 为 1。由于 _timer_running >= DEPTH (1 >= 1)，func1 的装饰器直接调用原始 func1 函数，跳过计时和打印。
   • func2 完成执行后，打印其计时信息，_timer_running 减至 0。 通过这种机制，只有最外层的函数调用（即 _timer_running 从 0 变为 1 的那次调用）才会触发计时和打印，内部嵌套的被装饰函数调用则会被静默处理。

灵活控制输出深度

这个解决方案的强大之处在于 DEPTH 参数的灵活性。你可以根据需要调整 DEPTH 的值，以控制哪些嵌套层级的函数调用应该打印其计时信息。

例如，如果我们将 DEPTH 设置为 2：

# 在 time_elapsed 装饰器内部，将 DEPTH 改为 2
# DEPTH = 2

再次运行代码，输出将变为：

--- Testing func1 ---
func1 took 0.10 seconds.

--- Testing func2 ---
func1 took 0.10 seconds.  # func2 内部调用的 func1 也被打印了
func2 took 0.30 seconds.

--- Testing func3 ---
func1 took 0.10 seconds.
func2 took 0.30 seconds.
func3 took 0.70 seconds.

--- Testing func4 ---
func1 took 0.10 seconds.
func2 took 0.30 seconds.
func3 took 0.70 seconds.
func4 took 1.50 seconds.

解释： 当 DEPTH = 2 时，_timer_running 在小于 2 的情况下会触发计时和打印。

• func2 内部调用 func1 时，_timer_running 从 0 变为 1 (为 func2 计时)，然后 func1 被调用。此时 _timer_running 为 1，小于 DEPTH (2)，因此 func1 的装饰器也会增加 _timer_running 到 2，执行计时和打印，然后减至 1。最后 func2 装饰器减至 0。
• 对于更深层次的嵌套，例如 func4 内部调用 func3，func3 内部调用 func2，func2 内部调用 func1：当 _timer_running 达到 2 或更高时，内部的装饰器将不再打印。例如，func4 计时时 _timer_running 为1，func3 计时时 _timer_running 为2，此时 func2 和 func1 的计时器将跳过打印。

注意事项与总结

1. 装饰器状态管理： 我们通过将 _timer_running 属性直接附加到 time_elapsed 函数对象上，实现了在所有被 @time_elapsed 装饰的函数实例之间共享一个状态。这种方式简单有效，但需要注意其作用域。
2. 线程安全： 如果你的应用程序是多线程的，并且多个线程可能同时调用被装饰的函数，那么 time_elapsed._timer_running 作为一个共享的可变状态，将存在竞态条件（race condition）。在这种情况下，你需要使用线程锁（如 threading.Lock）来保护 _timer_running 的读写操作，以确保线程安全。
3. 通用性： 这种基于深度计数的策略不仅适用于计时装饰器，也适用于任何需要在嵌套调用中控制输出或行为的装饰器场景。
4. 清晰的逻辑： try...finally 块的使用确保了 _timer_running 计数器无论函数执行是否发生异常，都能正确地递减，保持状态的准确性。

通过这种深度计数机制，我们成功地解决了Python装饰器在嵌套函数调用中产生的冗余输出问题，同时提供了灵活的控制能力，使得开发者可以根据实际需求调整输出的粒度。这是一种优雅且实用的装饰器设计模式，值得在日常开发中借鉴和应用。