如何用Python AST自动定位特定函数调用并提取相关信息？

核心优化：用ast.NodeVisitor自动遍历节点 #

别再手动用ast.iter_child_nodes遍历了，直接继承ast.NodeVisitor重写visit_Call方法，它会自动递归遍历所有AST节点，完全不需要预先知道节点位置。

具体实现步骤 #

• 识别目标函数：针对model.fit()这类属性调用，先判断Call节点的func是否为ast.Attribute类型，通过func.attr匹配方法名（比如fit），通过func.value拿到调用对象。
• 提取参数：直接从Call节点的args（位置参数）和keywords（关键字参数）中提取内容。

示例代码 #

import ast

class TargetCallVisitor(ast.NodeVisitor):
    def __init__(self, target_method):
        self.target_method = target_method
        self.matches = []

    def visit_Call(self, node):
        # 过滤出属性调用（如obj.method()）
        if isinstance(node.func, ast.Attribute):
            if node.func.attr == self.target_method:
                # 解析调用对象名称
                caller = self._get_node_identifier(node.func.value)
                # 提取位置参数
                pos_args = [self._get_node_identifier(arg) for arg in node.args]
                # 提取关键字参数
                kw_args = {kw.arg: self._get_node_identifier(kw.value) for kw in node.keywords}
                
                self.matches.append({
                    'line': node.lineno,
                    'caller': caller,
                    'pos_args': pos_args,
                    'kw_args': kw_args
                })
        # 继续遍历子节点，不遗漏嵌套调用
        self.generic_visit(node)

    def _get_node_identifier(self, node):
        # 处理普通变量名（如model、X_train）
        if isinstance(node, ast.Name):
            return node.id
        # 如需处理更复杂的表达式（如obj.sub_obj），可扩展此处逻辑
        return ast.dump(node, indent=2)

# 测试用例
test_code = """
model = build_model()
model.fit(X_train, y_train, epochs=10, batch_size=32)
backup_model.fit(X_val, y_val)
"""

# 解析代码并匹配目标方法
tree = ast.parse(test_code)
visitor = TargetCallVisitor(target_method='fit')
visitor.visit(tree)

# 输出结果
for match in visitor.matches:
    print(f"行号: {match['line']}")
    print(f"调用对象: {match['caller']}")
    print(f"位置参数: {match['pos_args']}")
    print(f"关键字参数: {match['kw_args']}")
    print("---")

关键细节说明 #

• ast.NodeVisitor的generic_visit方法会自动递归遍历所有子节点，确保不会遗漏嵌套在其他结构里的调用。
• _get_node_identifier方法可根据需求扩展，比如支持解析ast.Attribute类型的调用对象（如parent_model.child.fit），或者处理函数返回值作为调用对象的情况（如get_model().fit）。
• 通过node.lineno可以直接获取调用所在的行号，快速定位代码位置。

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