added write, read and for statements

improved/ast.py

@@ -388,6 +388,42 @@ class IfStatement(Statement):
             if self.on_false:
                 self.on_false.eval(env)
 
+class ReadStatement(Statement):
+    def __repr__(self):
+        return 'ReadStatement()'
+
+    def eval(self, env):
+        try:
+            return int(input())
+        except ValueError:
+            return 0
+
+class WriteStatement(Statement):
+    """The AST node for write statements"""
+    def __init__(self, message):
+        """
+            The initialization method.
+
+            message -- the message that should be printed
+        """
+        self.message = message
+
+    def __repr__(self):
+        """
+            A representation of the node for debug purposes.
+
+            returns -- a String of the form
+                       "WriteStatement(message)"
+        """
+        return 'WriteStatement(%s)' % message
+
+    def eval(self, env):
+        """
+            Evaluates the node. Prints the message.
+
+            returns -- None
+        """
+        return print(self.message.eval(env))
 
 class WhileStatement(Statement):
     """The AST node for while statements"""
@@ -420,3 +456,34 @@ class WhileStatement(Statement):
         """
         while self.condition.eval(env):
             self.body.eval(env)
+
+class ForStatement(Statement):
+    """The AST node for for statements"""
+    def __init__(self, count, body):
+        """
+            The initialization method.
+
+            count -- the number of times the body should be executed
+            body  -- the statement that is to be executed while
+                     inside the loop
+        """
+        self.count = count
+        self.body = body
+
+    def __repr__(self):
+        """
+            A representation of the node for debug purposes.
+
+            returns -- a String of the form
+                       "ForStatement(count, body)"
+        """
+        return 'ForStatement(%s, %s)' % (self.count, self.body)
+
+    def eval(self, env):
+        """
+            Evaluates the node. Executes the body block count times.
+
+            env -- the environment in which to evaluate the node
+        """
+        for _ in range(self.count.eval(env)):
+            self.body.eval(env)