*动态*装饰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 @-synthax

@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

Answer 1

要做你想做的事，你的

traverse_tree

方法必须知道它是如何被调用的，并且没有简单的方法可以做到这一点。你可以做这样的事情...

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


root = generate_tree()
_traverse = with_recursion_depth(traverse_tree)
for item in _traverse(root, func=_traverse):
    print(item)

...但这有点像黑客。