如何使用Python的ast库检测函数内所有数学运算及对应操作数？

你之前的实现是静态解析常量表达式求值，没办法处理函数里的变量参数，所以采用AST插桩修改函数代码的方案更适配需求，核心逻辑是运行时自动输出每步运算的细节，整体步骤如下：

• 用ast.parse()的mode='exec'模式解析目标函数的完整源代码，得到完整的函数定义AST
• 自定义AST节点转换器，遍历所有ast.BinOp二元运算节点，把原生运算替换成自定义的带日志输出的运算函数
• 修复AST的节点位置信息，编译修改后的AST为可执行代码，加载到当前命名空间后直接调用即可自动输出运算日志

import ast
import inspect

# 自定义带日志的运算包装函数
def traced_op(left, op, right):
    op_name = op.__class__.__name__
    print(f"运算类型: {op_name}")
    print(f"左操作数: {left}")
    print(f"右操作数: {right}")
    if isinstance(op, ast.Add):
        res = left + right
    elif isinstance(op, ast.Sub):
        res = left - right
    elif isinstance(op, ast.Mult):
        res = left * right
    elif isinstance(op, ast.Div):
        res = left / right
    elif isinstance(op, ast.Pow):
        res = left ** right
    else:
        raise ValueError(f"不支持的运算符: {op_name}")
    print(f"运算结果: {res}\n")
    return res

# AST节点转换器，替换所有二元运算为插桩函数
class BinOpTracer(ast.NodeTransformer):
    def visit_BinOp(self, node):
        # 先递归处理子节点，保证运算顺序和原函数一致
        self.generic_visit(node)
        # 替换原BinOp节点为自定义运算调用
        return ast.Call(
            func=ast.Name(id='traced_op', ctx=ast.Load()),
            args=[node.left, ast.Constant(value=node.op), node.right],
            keywords=[]
        )

# 函数插桩工具
def trace_function(func):
    # 获取目标函数的源代码
    src = inspect.getsource(func)
    # 解析为AST
    tree = ast.parse(src)
    # 转换AST
    modified_tree = BinOpTracer().visit(tree)
    # 补全缺失的节点位置信息，避免编译报错
    ast.fix_missing_locations(modified_tree)
    # 编译并加载插桩后的函数到全局命名空间
    exec(compile(modified_tree, filename='<ast>', mode='exec'), globals())
    return globals()[func.__name__]

# 你的目标函数
def foo(x, y=2):
    z = x*y + 3
    print("z is ", z)
    return z**2

# 测试调用
traced_foo = trace_function(foo)
print("最终返回值:", traced_foo(3,2))

运算类型: Mult
左操作数: 3
右操作数: 2
运算结果: 6

运算类型: Add
左操作数: 6
右操作数: 3
运算结果: 9

z is  9
运算类型: Pow
左操作数: 9
右操作数: 2
运算结果: 81

最终返回值: 81

• 该代码兼容Python 3.8+，刚好匹配你使用的3.8.8版本，如果你需要适配更早Python版本，把ast.Constant替换为对应版本的ast.Num即可
• 目前只实现了二元运算的插桩，如果需要处理一元运算（比如-x）、比较运算，新增对应ast.UnaryOp、ast.Compare节点的访问逻辑即可
• 如果函数引用了外部作用域的变量，把对应的变量传入exec的第二个globals参数即可避免找不到变量的问题

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