如何用Python ast实现预解析函数字典表达式的装饰器？

问题描述 #

我想写一个装饰器，能预解析函数里的部分表达式，把指定字典中的字符串表达式替换成对应数值。示例如下：

原代码：

CONSTS = {'pos': 1}

@resolve(CONSTS)
def f(x: list):
    return x[CONSTS['pos']]

装饰后要等价于：

def f(x: list):
    return x[1]

这个装饰器只处理单个字典的字符串表达式替换，因为后续还要把转换后的函数传给另一个AST处理的装饰器，而那个装饰器不支持字典使用。请问能用AST模块实现吗？

解决方案：用AST模块完全可行 #

AST（抽象语法树）模块专门用于解析、修改Python代码的语法结构，正好适配这种静态代码替换的需求。以下是具体实现思路和代码：

核心实现逻辑 #

1. 解析被装饰函数的源码为AST节点树
2. 遍历AST树，精准识别字典名['键名']形式的节点：
   • 匹配*Subscript*节点（对应下标访问语法）
   • 其value属性为*Name*节点，且id与传入的字典名称一致
   • 其slice属性为*Constant*节点，且值为存在于目标字典中的字符串键
3. 将符合条件的节点直接替换为字典中对应的常量值
4. 把修改后的AST重新编译为函数对象，替换原函数

完整代码实现 #

import ast
import inspect
from functools import wraps

def resolve(consts_dict):
    # 获取传入字典的变量名（比如传入CONSTS时，拿到字符串"CONSTS"）
    frame = inspect.currentframe().f_back
    consts_name = next(k for k, v in frame.f_locals.items() if v is consts_dict)

    class ConstReplacer(ast.NodeTransformer):
        def visit_Subscript(self, node):
            # 检查是否是目标字典的键访问形式
            if (isinstance(node.value, ast.Name) 
                and node.value.id == consts_name
                and isinstance(node.slice, ast.Constant)
                and isinstance(node.slice.value, str)
                and node.slice.value in consts_dict):
                # 替换为对应的常量值
                return ast.Constant(value=consts_dict[node.slice.value])
            # 非目标节点则继续遍历子节点
            return self.generic_visit(node)

    def decorator(func):
        # 获取原函数的源码
        source = inspect.getsource(func)
        # 解析为AST树
        tree = ast.parse(source)
        # 执行节点替换
        replacer = ConstReplacer()
        modified_tree = replacer.visit(tree)
        # 修复AST的位置信息，避免编译报错
        ast.fix_missing_locations(modified_tree)
        
        # 编译修改后的AST，生成新函数
        func_globals = func.__globals__.copy()
        exec(compile(modified_tree, filename=inspect.getfile(func), mode='exec'), func_globals)
        new_func = func_globals[func.__name__]
        
        # 保留原函数的元信息（文档字符串、参数注解等）
        wraps(func)(new_func)
        return new_func

    return decorator

测试验证 #

用示例代码测试：

CONSTS = {'pos': 1}

@resolve(CONSTS)
def f(x: list):
    return x[CONSTS['pos']]

# 调用测试，输出20
print(f([10, 20]))
# 查看修改后的函数源码，可看到x[1]的替换结果
import inspect
print(inspect.getsource(f))

注意事项 #

• 该实现假设字典在装饰器所在的作用域中定义，若字典是导入的，需要调整获取字典名称的逻辑
• 仅支持直接的字典名['键名']形式，不处理嵌套字典或更复杂的表达式
• 替换后会完整保留原函数的元信息，确保后续装饰器能正常处理

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