Python中如何动态装饰递归函数且不修改当前作用域

问题背景 #

需要为递归遍历二叉树的函数traverse_tree动态添加递归深度统计功能，核心要求：

• 动态实现装饰，不修改当前作用域中的原函数
• 内部递归调用也需被装饰，正确返回递归深度
• 保留未装饰的原函数，按需使用带深度统计的版本

使用@with_recursion_depth语法装饰时，递归调用能正确被装饰，输出每个节点值对应的深度；但动态调用with_recursion_depth(traverse_tree)(root)时，仅外层调用被装饰，内部递归调用仍指向原函数，导致所有深度都显示为0。直接赋值traverse_tree = with_recursion_depth(traverse_tree)虽能解决，但会覆盖原函数，不符合需求。

现有代码示例 #

基础定义与装饰器 #

import functools

class Node:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

def generate_tree():
    root = Node(1)
    root.left = Node(2)
    root.right = Node(3)
    root.left.left = Node(4)
    root.left.right = Node(5)
    root.right.left = Node(6)
    root.right.right = Node(7)
    return root

def with_recursion_depth(func):
    """Yield recursion depth alongside original values of an iterator."""
    
    class Depth(int): pass
    depth = Depth(-1)

    def depth_in_value(value, depth) -> bool:
        return isinstance(value, tuple) and len(value) == 2 and value[-1] is depth

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        nonlocal depth
        depth = Depth(depth + 1)
        for value in func(*args, **kwargs):
            if depth_in_value(value, depth):
                yield value
            else:
                yield value, depth
        depth = Depth(depth - 1)

    return wrapper

装饰器语法（正常工作） #

@with_recursion_depth
def traverse_tree(node):
    """Recursively yield values of the binary tree."""
    yield node.value
    if node.left:
        yield from traverse_tree(node.left)
    if node.right:
        yield from traverse_tree(node.right)

root = generate_tree()
for item in traverse_tree(root):
    print(item)
# 输出:
# (1, 0)
# (2, 1)
# (4, 2)
# (5, 2)
# (3, 1)
# (6, 2)
# (7, 2)

动态调用（问题场景） #

def traverse_tree(node):
    """Recursively yield values of the binary tree."""
    yield node.value
    if node.left:
        yield from traverse_tree(node.left)
    if node.right:
        yield from traverse_tree(node.right)

root = generate_tree()
for item in with_recursion_depth(traverse_tree)(root):
    print(item)
# 输出:
# (1, 0)
# (2, 0)
# (4, 0)
# (5, 0)
# (3, 0)
# (6, 0)
# (7, 0)

解决方案 #

方案1：添加可选参数传递装饰后的函数（最简单） #

修改原函数，增加一个可选参数self_func，让递归调用使用该参数指向的函数，而非硬编码的原函数名。这样装饰后的wrapper可以将自身传递给原函数，确保内部递归调用也使用装饰后的版本。

def traverse_tree(node, self_func=None):
    # 若未传入self_func，默认指向原函数
    self_func = self_func or traverse_tree
    yield node.value
    if node.left:
        yield from self_func(node.left, self_func)
    if node.right:
        yield from self_func(node.right, self_func)

# 动态装饰并调用，原函数traverse_tree仍保留
decorated_traverse = with_recursion_depth(traverse_tree)
for item in decorated_traverse(root):
    print(item)
# 输出与@语法装饰的结果一致

方案2：修复AST修改方案 #

原AST修改代码的问题在于每次递归调用都会重新装饰函数，导致状态混乱。修复后的方案会在函数内部创建一个指向装饰后函数的局部变量，替换所有递归调用的目标。

import ast
import inspect
import types

def modify_recursive_calls(func, decorator):
    # 解析原函数源码为AST
    source = inspect.getsource(func)
    tree = ast.parse(source)
    
    func_def = tree.body[0]
    # 在函数开头添加局部变量，指向装饰后的原函数
    decorated_assign = ast.Assign(
        targets=[ast.Name(id="_decorated_self", ctx=ast.Store())],
        value=ast.Call(
            func=ast.Name(id=decorator.__name__, ctx=ast.Load()),
            args=[ast.Name(id=func.__name__, ctx=ast.Load())],
            keywords=[]
        )
    )
    func_def.body.insert(0, decorated_assign)
    
    # 替换所有递归调用为_decorated_self
    def replace_recursive_calls(node):
        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == func.__name__:
            node.func = ast.copy_location(ast.Name(id="_decorated_self", ctx=ast.Load()), node.func)
        # 递归处理所有子节点
        for field, value in ast.iter_fields(node):
            if isinstance(value, list):
                for item in value:
                    if isinstance(item, ast.AST):
                        replace_recursive_calls(item)
            elif isinstance(value, ast.AST):
                replace_recursive_calls(item)
    
    replace_recursive_calls(func_def)
    ast.fix_missing_locations(tree)
    
    # 编译AST并生成新函数
    compiled = compile(tree, filename="<ast_modified>", mode="exec")
    new_func = types.FunctionType(
        compiled.co_consts[1], 
        func.__globals__, 
        func.__name__, 
        func.__defaults__, 
        func.__closure__
    )
    return new_func

# 使用方式
def traverse_tree(node):
    yield node.value
    if node.left:
        yield from traverse_tree(node.left)
    if node.right:
        yield from traverse_tree(node.right)

# 生成修改后的装饰函数，原函数不受影响
decorated_traverse = modify_recursive_calls(traverse_tree, with_recursion_depth)
for item in decorated_traverse(root):
    print(item)
# 输出与@语法装饰的结果一致

总结 #

• 方案1实现最简单，仅需修改原函数添加一个可选参数，无需复杂的AST或字节码操作，推荐优先使用。
• 方案2无需修改原函数，通过AST修改生成新的装饰后函数，适合无法修改原函数源码的场景。

内容的提问来源于stack exchange，提问作者Michiel Karrenbelt