1 files changed, 50 insertions, 22 deletions

diff --git a/src/ast/Python3VisitorImpl.java b/src/ast/Python3VisitorImpl.java
index 604c8d2..d4b4fbf 100644
--- a/src/ast/Python3VisitorImpl.java
+++ b/src/ast/Python3VisitorImpl.java
@@ -21,17 +21,19 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
      * ```
      */
     public Node visitRoot(RootContext ctx) {
-        ArrayList<Node> stmts = new ArrayList<Node>();
-        ArrayList<Node> compStmts = new ArrayList<Node>();
-
-        for (Simple_stmtsContext stm : ctx.simple_stmts()) {
-            stmts.add(visit(stm));
-        }
-        for (Compound_stmtContext stm : ctx.compound_stmt()) {
-            compStmts.add(visit(stm));
+        ArrayList<Node> childs = new ArrayList<Node>();
+        
+        for (int i = 0; i < ctx.getChildCount(); i++){
+            var child = ctx.getChild(i);
+            
+            if (child instanceof Simple_stmtsContext) {
+                childs.add(visit((Simple_stmtsContext) child));
+            } else if (child instanceof Compound_stmtContext) {
+                childs.add(visit((Compound_stmtContext) child));
+            }
         }
 
-        return new RootNode(stmts, compStmts);
+        return new RootNode(childs);
     }
 
     /**
@@ -262,10 +264,28 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
             x = new AugassignNode(ctx.SUB_ASSIGN());
         } else if (ctx.MULT_ASSIGN() != null) {
             x = new AugassignNode(ctx.MULT_ASSIGN());
+        } else if (ctx.AT_ASSIGN() != null) {
+            x = new AugassignNode(ctx.AT_ASSIGN());
         } else if (ctx.DIV_ASSIGN() != null) {
             x = new AugassignNode(ctx.DIV_ASSIGN());
-        }
-
+        } else if (ctx.MOD_ASSIGN() != null) {
+            x = new AugassignNode(ctx.MOD_ASSIGN());
+        } else if (ctx.AND_ASSIGN() != null) {
+            x = new AugassignNode(ctx.AND_ASSIGN());
+        } else if (ctx.OR_ASSIGN() != null) {
+            x = new AugassignNode(ctx.OR_ASSIGN());
+        } else if (ctx.XOR_ASSIGN() != null) {
+            x = new AugassignNode(ctx.XOR_ASSIGN());
+        } else if (ctx.LEFT_SHIFT_ASSIGN() != null) {
+            x = new AugassignNode(ctx.LEFT_SHIFT_ASSIGN());
+        } else if (ctx.RIGHT_SHIFT_ASSIGN() != null) {
+            x = new AugassignNode(ctx.RIGHT_SHIFT_ASSIGN());
+        } else if (ctx.POWER_ASSIGN() != null) {
+            x = new AugassignNode(ctx.POWER_ASSIGN());
+        } else if (ctx.IDIV_ASSIGN() != null) {
+            x = new AugassignNode(ctx.IDIV_ASSIGN());
+        }
+        
         return x;
     }
 
@@ -331,17 +351,19 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
      * ```
      */
     public Node visitBlock(BlockContext ctx) {
-        ArrayList<Node> stmts = new ArrayList<Node>();
-        ArrayList<Node> compStmts = new ArrayList<Node>();
-
-        for (Simple_stmtsContext s : ctx.simple_stmts()) {
-            stmts.add(visit(s));
-        }
-        for (Compound_stmtContext s : ctx.compound_stmt()) {
-            compStmts.add(visit(s));
+        ArrayList<Node> childs = new ArrayList<Node>();
+
+        for (int i = 0; i < ctx.getChildCount(); i++){
+            var child = ctx.getChild(i);
+            
+            if (child instanceof Simple_stmtsContext) {
+                childs.add(visit((Simple_stmtsContext) child));
+            } else if (child instanceof Compound_stmtContext) {
+                childs.add(visit((Compound_stmtContext) child));
+            }
         }
 
-        return new BlockNode(stmts, compStmts);
+        return new BlockNode(childs);
     }
 
     /**
@@ -491,7 +513,7 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
             methodCall = ctx.NAME();
         }
 
-        return new TrailerNode(arglist, exprs, methodCall);
+        return new TrailerNode(arglist, exprs, methodCall, ctx.OPEN_PAREN() != null);
     }
 
     /**
@@ -503,7 +525,13 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
      * ```
      */
     public Node visitExprlist(ExprlistContext ctx) {
-        Node exp = visit(ctx.expr(0));
+        Node exp;
+
+        if (ctx.expr(0).expr(0) != null){
+            exp = visit(ctx.expr(0).expr(0));
+        } else {
+            exp = visit(ctx.expr(0));
+        }
 
         return exp;
     }