在Python中动态装饰递归函数

Question

我遇到了一个情况，需要在Python中动态地给一个函数的递归调用添加装饰。关键要求是要做到这一点，而不修改当前作用域中的函数。让我来解释一下这个情况以及我目前尝试过的。

假设我有一个函数 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

# 1. using @-syntax
@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)

# Output:
# (1, 0)
# (2, 1)
# (4, 2)
# (5, 2)
# (3, 1)
# (6, 2)
# (7, 2)


# 2. Dynamically:  
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)

# Output:
# (1, 0)
# (2, 0)
# (4, 0)
# (5, 0)
# (3, 0)
# (6, 0)
# (7, 0)

看起来问题出在如何给函数内部的递归调用添加装饰。当我动态使用装饰器时，它只装饰了外部函数的调用，而没有装饰函数内部的递归调用。我可以通过重新赋值来实现这个功能（traverse_tree = with_recursion_depth(traverse_tree)），但这样一来，函数在当前作用域中就被修改了。我希望能够动态地实现这个，这样我就可以选择使用未装饰的函数，或者选择性地包装它以获取递归深度的信息。

我希望保持事情简单，尽量避免像字节码操作这样的复杂技术，如果有其他解决方案的话。不过，如果这是必须的路径，我也愿意去探索。我在这方面尝试过，但还没有成功。

import ast


def modify_recursive_calls(func, decorator):

    def decorate_recursive_calls(node):
        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == func.__name__:
            func_name = ast.copy_location(ast.Name(id=node.func.id, ctx=ast.Load()), node.func)
            decorated_func = ast.Call(
                func=ast.Name(id=decorator.__name__, ctx=ast.Load()),
                args=[func_name],
                keywords=[],
            )
            node.func = decorated_func
        for field, value in ast.iter_fields(node):
            if isinstance(value, list):
                for item in value:
                    if isinstance(item, ast.AST):
                        decorate_recursive_calls(item)
            elif isinstance(value, ast.AST):
                decorate_recursive_calls(value)

    tree = ast.parse(inspect.getsource(func))
    decorate_recursive_calls(tree)
    ast.fix_missing_locations(tree)
    modified_code = compile(tree, filename="<ast>", mode="exec")
    modified_function = types.FunctionType(modified_code.co_consts[1], func.__globals__)
    return modified_function

代码优化递归装饰器动态编程二叉树递归深度函数包装作用域管理

在Python中动态装饰递归函数

3 个回答

撰写回答