1 files changed, 28 insertions, 4 deletions

diff --git a/src/ast/nodes/AtomNode.java b/src/ast/nodes/AtomNode.java
index ea99808..204c3b3 100644
--- a/src/ast/nodes/AtomNode.java
+++ b/src/ast/nodes/AtomNode.java
@@ -13,6 +13,8 @@ import semanticanalysis.SymbolTable;
  */
 public class AtomNode implements Node {
 
+    private String prefix;
+    private String suffix;
     protected String val;
     protected TestlistCompNode exprlist;
 
@@ -21,8 +23,14 @@ public class AtomNode implements Node {
     protected int offset;
 
     public AtomNode(String val, Node exprlist) {
+        this(val, exprlist, "", "");
+    }
+
+    public AtomNode(String val, Node exprlist, String prefix, String suffix) {
         this.val = val;
         this.exprlist = (TestlistCompNode) exprlist;
+        this.prefix = prefix;
+        this.suffix = suffix;
     }
 
     /**
@@ -33,6 +41,10 @@ public class AtomNode implements Node {
         return val;
     }
 
+    public void setId(String id) {
+        this.val = id;
+    }
+
     @Override
     public ArrayList<SemanticError> checkSemantics(SymbolTable ST, int _nesting, FunctionType ft) {
         ArrayList<SemanticError> errors = new ArrayList<>();
@@ -134,13 +146,25 @@ public class AtomNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
-        // FIXME: can be a testlist_comp with two expr and two atoms
+    public String printAST(String prefix) {
         if (val != null) {
-            return prefix + "Atom(" + val + ")\n";
+            return prefix + "AtomNode: " + val + "\n";
+        } else {
+            return prefix + "AtomNode\n" + exprlist.printAST(prefix + "  ");
         }
+    }
 
-        return prefix + "Atom(null)\n";
+    @Override
+    public String toPrint(String prefix) {
+        String str = prefix + this.prefix;
+
+        if (val != null) {
+            str += val;
+        } else {
+            str += this.exprlist.toPrint("");
+        }
 
+        str += this.suffix;
+        return str;
     }
 }