如何使用ANTLR为Python中每个函数生成独立AST

嘿，看你已经有Python语法基础，现在想用ANTLR的Java实现来提取每个Python函数的独立AST对吧？我看了你给出的初始代码，帮你补全并优化了完整实现，直接就能用：

• 继承Python3BaseListener，通过监听函数的进入/退出事件，精准捕获每个函数对应的语法节点
• 利用ANTLR的ParseTreeAPI，把每个函数对应的子树单独提取出来，作为独立AST
• 内置了AST打印逻辑，方便你直观查看每个函数的结构，也可以改成保存到文件的逻辑

import org.antlr.v4.runtime.*;
import org.antlr.v4.runtime.tree.ParseTree;
import org.antlr.v4.runtime.tree.TerminalNode;
import java.io.IOException;
import java.io.StringReader;

public class TestGrammar extends Python3BaseListener {
    // 临时存储当前函数的AST根节点
    private ParseTree currentFunctionAST;
    // 标记是否处于函数定义内部，避免混淆嵌套函数
    private boolean inFunction = false;

    public static void main(String[] args) throws IOException {
        // 读取目标Python源代码文件
        PlagiarismPercentage obj = new PlagiarismPercentage();
        String source = obj.readFile("C:\\Users\\Paridhi\\quickSort.py");
        
        // 初始化ANTLR词法/语法分析流程
        ANTLRInputStream inputCharStream = new ANTLRInputStream(new StringReader(source));
        Python3Lexer lexer = new Python3Lexer(inputCharStream);
        CommonTokenStream tokenStream = new CommonTokenStream(lexer);
        Python3Parser parser = new Python3Parser(tokenStream);
        
        // 获取整个Python文件的语法树根节点
        ParseTree root = parser.file_input();
        
        // 创建自定义Listener并遍历语法树
        TestGrammar functionAstExtractor = new TestGrammar();
        ParseTreeWalker.DEFAULT.walk(functionAstExtractor, root);
    }

    // 进入函数定义节点时触发，开始提取AST
    @Override
    public void enterFuncdef(Python3Parser.FuncdefContext ctx) {
        inFunction = true;
        // 记录当前函数的AST根节点（就是funcdef本身）
        currentFunctionAST = ctx;
        System.out.println("=== 提取到新函数的AST ===");
        // 打印AST结构，也可以替换成保存到文件的逻辑
        printAST(currentFunctionAST, 0);
    }

    // 退出函数定义节点时触发，重置状态
    @Override
    public void exitFuncdef(Python3Parser.FuncdefContext ctx) {
        inFunction = false;
        currentFunctionAST = null;
        System.out.println("\n=== 该函数AST提取完成 ===");
    }

    // 递归打印AST结构的辅助方法，带缩进更清晰
    private void printAST(ParseTree node, int indent) {
        // 打印缩进空格
        for (int i = 0; i < indent; i++) {
            System.out.print("  ");
        }
        
        // 处理终端节点（比如函数名、关键字、字面量）
        if (node instanceof TerminalNode) {
            TerminalNode terminal = (TerminalNode) node;
            System.out.println("终端节点: " + terminal.getText() + " (类型: " + Python3Lexer.VOCABULARY.getSymbolicName(terminal.getSymbol().getType()) + ")");
            return;
        }
        
        // 处理非终端节点（比如funcdef、parameter_list、suite等）
        ParserRuleContext ctx = (ParserRuleContext) node;
        System.out.println("非终端节点: " + Python3Parser.ruleNames[ctx.getRuleIndex()]);
        
        // 递归打印所有子节点
        for (int i = 0; i < ctx.getChildCount(); i++) {
            printAST(ctx.getChild(i), indent + 1);
        }
    }
}

// 给你补了个PlagiarismPercentage类的参考实现（如果你的已有实现不同可以忽略）
class PlagiarismPercentage {
    public String readFile(String filePath) throws IOException {
        return new String(java.nio.file.Files.readAllBytes(java.nio.file.Paths.get(filePath)));
    }
}

• enterFuncdef方法：这是ANTLR遍历到函数定义时的回调，我们在这里直接把当前函数的语法节点作为AST根节点，后续可以对这个节点做任何操作（比如序列化、分析）
• printAST方法：递归遍历AST的每个节点，区分终端/非终端节点并带缩进打印，能让你清晰看到函数的结构层次
• inFunction标记：如果你的Python代码里有嵌套函数，这个标记可以帮你区分当前处理的是外层还是内层函数的节点

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