Finish refactoring

Project-02/syntax.py

@@ -9,17 +9,6 @@ class EXPRESSION:
     def copy():
         return EXPRESSION()
 
-    def allNodes(self):
-        ret = [self]
-        for node in (self.__getattribute__(a) for a in self.__dict__.keys()):
-            if isinstance(node, EXPRESSION):
-                ret += node.allNodes()
-            if isinstance(node, list):
-                for n in node:
-                    if isinstance(n, EXPRESSION):
-                        ret += n.allNodes()
-        return ret
-
     def children(self):
         """Return a list of (name, childNode)."""
         out = []
@@ -31,33 +20,55 @@ class EXPRESSION:
             elif isinstance(value, list):
                 for i, elem in enumerate(value):
                     if isinstance(elem, EXPRESSION):
-                        out.append((f"{key}[{i}]", elem))
+                        out.append((f"{key}{i}", elem))
         return out
 
-    def to_dot(self, out):
-        # node label is class name or class name + value
+    def to_dot(self, visited=None, root=True):
+        if visited is None:
+            visited = set()
+
+        # Prevent infinite recursion
+        if id(self) in visited:
+            return ""
+        visited.add(id(self))
+
+        parts = []
+
+        # Add a header at the root node
+        if root:
+            parts.append("digraph AST {\n")
+
+        # Append to label
         label = type(self).__name__
-        if hasattr(self, "operator"):  # AOP/EQOP/COMP/LOP
+        if hasattr(self, "operator"):
             label += f"({self.operator})"
-        if hasattr(self, "name"):  # ID
+        if hasattr(self, "name"):
             label += f"({self.name})"
-        if hasattr(self, "value"):  # CONST
+        if hasattr(self, "value"):
             label += f"({self.value})"
 
-        out.write(f'  node{self.pp} [label="{label}"];\n')
+        parts.append(f'  node{self.pp} [label="{label}"];\n')
+
+        # Draw edges
+        for edge_name, child in self.children():
+            parts.append(f'  node{self.pp} -> node{child.pp} [label="{edge_name}"];\n')
+            parts.append(child.to_dot(visited, root=False))
+
+        # Add footer at the root node
+        if root:
+            parts.append("}\n")
+
+        return "".join(parts)
 
-        for (edge_name, child) in self.children():
-            out.write(f'  node{self.pp} -> node{child.pp} [label="{edge_name}"];\n')
-            child.to_dot(out)
 
 class LET(EXPRESSION):
-    def __init__(self, declarations, body):
+    def __init__(self, decls, body):
         super().__init__()
-        self.declarations = declarations
+        self.decl = decls
         self.body = body
 
     def __str__(self):
-        return "let " + ", ".join(str(d) for d in self.declarations) + " in " + str(self.body)
+        return "let " + ", ".join(str(d) for d in self.decl) + " in " + str(self.body)
 
 class DECL(EXPRESSION):
     def __init__(self, f_name, params, body):
@@ -70,13 +81,13 @@ class DECL(EXPRESSION):
         return f"{self.f_name}({self.params}) {{ {self.body} }}"
 
 class CALL(EXPRESSION):
-    def __init__(self, f_name, arguments):
+    def __init__(self, f_name, args):
         super().__init__()
         self.f_name = f_name
-        self.arguments = arguments
+        self.arg = args
 
     def __str__(self):
-        return self.f_name + "(" + ",".join(str(a) for a in self.arguments) + ")"
+        return self.f_name + "(" + ",".join(str(a) for a in self.arg) + ")"
 
 class ID(EXPRESSION):
     def __init__(self, name):
@@ -135,13 +146,13 @@ class LOP(EXPRESSION):
         return "(" + str(self.arg1) + " " + str(self.operator) + " " + str(self.arg2) + ")"
 
 class ASSIGN(EXPRESSION):
-    def __init__(self, variable, expression):
+    def __init__(self, var, expr):
         super().__init__()
-        self.variable = variable
-        self.expression = expression
+        self.var = var
+        self.expr = expr
 
     def __str__(self):
-        return self.variable.name + " = " + str(self.expression)
+        return self.var.name + " = " + str(self.expr)
 
 class SEQ(EXPRESSION):
     def __init__(self, exp1, exp2):
@@ -153,20 +164,20 @@ class SEQ(EXPRESSION):
         return str(self.exp1) + "; " + str(self.exp2)
 
 class IF(EXPRESSION):
-    def __init__(self, condition, exp1, exp2):
+    def __init__(self, cond, exp1, exp2):
         super().__init__()
-        self.condition = condition
+        self.cond = cond
         self.exp1 = exp1
         self.exp2 = exp2
 
     def __str__(self):
-        return "if (" + str(self.condition) + ") then { " + str(self.exp1) + " } else { " + str(self.exp2) + " }"
+        return "if (" + str(self.cond) + ") then { " + str(self.exp1) + " } else { " + str(self.exp2) + " }"
 
 class WHILE(EXPRESSION):
-    def __init__(self, condition, body):
+    def __init__(self, cond, body):
         super().__init__()
-        self.condition = condition
+        self.cond = cond
         self.body = body
 
     def __str__(self):
-        return "while (" + str(self.condition) + ") do { " + str(self.body) + " }"
+        return "while (" + str(self.cond) + ") do { " + str(self.body) + " }"