如何实现可删除函数内指定变量的Python装饰器？

原装饰器失效原因 #

你实现的del_variables装饰器无法生效，核心原因是：装饰器中wrapper函数的locals()指向的是自身的命名空间，和被装饰函数（比如foo）的局部命名空间完全独立，所以在finally块中操作的是wrapper的局部变量，根本触不到目标函数里的c等变量。

可行解决方案：通过AST修改原函数代码 #

要让装饰器能操作被装饰函数的局部变量，最可靠的方式是动态修改原函数的抽象语法树（AST），给原函数的代码包裹一层try/except/finally逻辑，让finally块直接属于被装饰函数的命名空间，自然就能访问到它的局部变量。

实现代码 #

import ast
import inspect
import functools

def del_variables(*variables):
    def decorator(func):
        # 获取原函数的源代码
        source = inspect.getsource(func)
        # 解析源代码为AST
        tree = ast.parse(source)
        # 定位到原函数的定义节点
        func_def = tree.body[0]
        
        # 构建finally块中的删除变量逻辑
        finally_body = []
        for var_name in variables:
            # 生成代码：if var_name in locals(): del locals()[var_name]
            in_check = ast.Compare(
                left=ast.Constant(value=var_name),
                ops=[ast.In()],
                comparators=[ast.Call(func=ast.Name(id="locals", ctx=ast.Load()), args=[], keywords=[])]
            )
            del_stmt = ast.Delete(
                targets=[ast.Subscript(
                    value=ast.Call(func=ast.Name(id="locals", ctx=ast.Load()), args=[], keywords=[]),
                    slice=ast.Constant(value=var_name),
                    ctx=ast.Del()
                )]
            )
            finally_body.append(ast.If(test=in_check, body=[del_stmt], orelse=[]))
        
        # 构建try节点，包裹原函数体
        try_node = ast.Try(
            body=func_def.body,
            handlers=[
                # 捕获Exception并重新抛出，保留原异常信息
                ast.ExceptHandler(
                    type=ast.Name(id="Exception", ctx=ast.Load()),
                    name="error",
                    body=[ast.Raise(exc=ast.Name(id="error", ctx=ast.Load()), cause=None)]
                )
            ],
            orelse=[],
            finalbody=finally_body
        )
        
        # 替换原函数体为带try/finally的逻辑
        func_def.body = [try_node]
        # 修复AST的位置信息，避免编译警告
        ast.fix_missing_locations(tree)
        
        # 编译修改后的AST
        code = compile(tree, filename=inspect.getsourcefile(func), mode="exec")
        # 在原函数的全局命名空间中执行编译后的代码
        ns = func.__globals__.copy()
        exec(code, ns)
        # 获取修改后的函数对象
        new_func = ns[func.__name__]
        # 复制原函数的元信息（如文档字符串、模块信息等）
        functools.update_wrapper(new_func, func)
        
        return new_func
    return decorator

使用示例 #

@del_variables("c")
def foo(a, b):
    c = 3
    print(f"函数内部c的值：{c}")
    return a + b + c

# 调用测试
result = foo(1, 2)
print(f"返回结果：{result}")

方案说明 #

1. 该装饰器通过inspect.getsource获取原函数的源代码，解析为AST后，自动给原函数体包裹try/except/finally逻辑。
2. finally块直接属于被装饰函数的命名空间，因此可以直接访问并删除指定的局部变量，触发对象的__del__方法。
3. 保留了原函数的所有元信息，使用方式完全透明，不需要修改原函数的代码。

其他注意事项 #

• 该方案依赖原函数的源代码可用性，如果函数是动态生成的（如lambda或exec创建的函数），可能无法正常工作，但测试场景下的函数一般都是静态定义的，完全适用。
• 如果你不需要捕获并重新抛出异常，可以删除try_node中的handlers列表，只保留body和finalbody。

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