from .astnodes import *

class Visitor:
    """ 
        Visitor that can visit all nodes in the AST
        and run pre and post functions.
    """
    def visit(self, node, f_pre=None, f_post=None):
        self.f_pre = f_pre
        self.f_post = f_post
        self.do(node)

    def do(self, node):
      # Run visitor:
      if self.f_pre:
            self.f_pre(node)

      # Descent into subnodes:
      if type(node) is Package:
            for decl in node.declarations:
                self.do(decl)
      elif type(node) is Function:
            for s in node.declarations:
                self.do(s)
            self.do(node.body)
      elif type(node) is CompoundStatement:
         for s in node.statements:
            self.do(s)
      elif type(node) is IfStatement:
         self.do(node.condition)
         self.do(node.truestatement)
         self.do(node.falsestatement)
      elif type(node) is FunctionCall:
         for arg in node.args:
            self.do(arg)
      elif type(node) is Assignment:
         self.do(node.lval)
         self.do(node.rval)
      elif type(node) is ReturnStatement:
         self.do(node.expr)
      elif type(node) is Binop:
         self.do(node.a)
         self.do(node.b)
      elif type(node) is Constant:
         self.do(node.value)
      elif type(node) in [EmptyStatement, VariableUse, Variable, Literal, FunctionType]:
         # Those nodes do not have child nodes.
         pass
      elif type(node) is WhileStatement:
         self.do(node.condition)
         self.do(node.statement)
      else:
           raise Exception('Could not visit "{0}"'.format(node))

      # run post function
      if self.f_post:
            self.f_post(node)