How to get all variables from ANTLR AST in java

Question

I have built AST using ANTLR for two simple assignments x = 1 + 2; y = 3 + 5;

and the tree is generated as follows:

single_input
 simple_stmt
   small_stmt
     expr_stmt
       testlist_star_expr
         test
           or_test
             and_test
               not_test
                 comparison
                   expr
                     xor_expr
                       and_expr
                         shift_expr
                           arith_expr
                             term
                               factor
                                 power
                                   atom_expr
                                     atom
   testlist_star_expr
     test
       or_test
         and_test
           not_test
             comparison
               expr
                 xor_expr
                   and_expr
                     shift_expr
                       arith_expr
                         term
                           factor
                             power
                               atom_expr
                                 atom
                         term
                           factor
                             power
                               atom_expr
                                 atom
small_stmt
 expr_stmt
   testlist_star_expr
     test
       or_test
         and_test
           not_test
             comparison
               expr
                 xor_expr
                   and_expr
                     shift_expr
                       arith_expr
                         term
                           factor
                             power
                               atom_expr
                                 atom
   testlist_star_expr
     test
       or_test
         and_test
           not_test
             comparison
               expr
                 xor_expr
                   and_expr
                     shift_expr
                       arith_expr
                         term
                           factor
                             power
                               atom_expr
                                 atom
                         term
                           factor
                             power
                               atom_expr
                                 atom

I want to know how to access the variable x in the atom expression. How do I parse the tree?

I have printed the above tree as follows:

RuleContext r = parser.single_input();
ASTPrinter ast = new ASTPrinter();
ast.print(r);

And my ASTPrinter class is as follows:

public class ASTPrinter {

public void print(RuleContext ctx) {
   explore(ctx, 0);
}

private void explore(RuleContext ctx, int indentation) {
   String ruleName = Python3Parser.ruleNames[ctx.getRuleIndex()];
   for (int i=0;i<indentation;i++) {
       System.out.print("  ");
   }
   System.out.println(ruleName);
   for (int i=0;i<ctx.getChildCount();i++) {
       ParseTree element = ctx.getChild(i);
       if (element instanceof RuleContext) {
           explore((RuleContext)element, indentation + 1);
       }
   }
  }

  }

I can get the ruleName from the ruleNames array. Is there any way I can retrieve the name of my variable(x)?

Answer 1

Assuming you're using the Pyhton3.g4 grammar , you can display the variables like this:

public class Main {

    public static void main(String[] args) {

        String source = "x = 1 + 2; y = 3 + 5;\n";
        Python3Lexer lexer = new Python3Lexer(CharStreams.fromString(source));
        Python3Parser parser = new Python3Parser(new CommonTokenStream(lexer));

        ParseTreeWalker.DEFAULT.walk(new Python3BaseListener() {

            // expr_stmt
            //  : testlist_star_expr ( augassign ( yield_expr | testlist)
            //                       | ( '=' ( yield_expr| testlist_star_expr ) )*
            //                       )
            //  ;
            @Override
            public void enterExpr_stmt(Python3Parser.Expr_stmtContext ctx) {
                System.out.println("variable: " + ctx.testlist_star_expr().get(0).getText());
            }
        }, parser.single_input());
    }
}

which will print:

variable: x
variable: y

If you want to know more about listeners and how to "decorate" parser rules with labels so you can fine-tune tree-walking, have a look at these links:

• https://github.com/antlr/antlr4/blob/master/doc/parser-rules.md
• https://github.com/antlr/antlr4/blob/master/doc/listeners.md