modelverse/interface/HUTN/hutn_compiler/prettyprint_visitor.py

from hutnparser import Tree
from visitor import Visitor

class PrintVisitor(Visitor):
    def __init__(self, args):
        Visitor.__init__(self, args)
        self.result = ''
        pass

    def visit(self, tree, level=0):
        if isinstance(tree, Tree):
            self.result += '\n' + '\t' * level + tree.head
            for child in tree.tail:
                self.visit(child, level+1)
        else:
            self.result += '(' + repr(tree) + ')'

    def dump(self):
        return self.result

class PrettyPrintVisitor(Visitor):
    def __init__(self, args):
        Visitor.__init__(self, args)
        self.result = ""

    def visit(self, tree, level=0):
        if isinstance(tree, Tree):
            if tree.is_rule() and len(tree.get_tail()) > 1:
                self.result += '\n' + '\t' * level + tree.head
                if not tree.head.startswith("implicit_autogenerated_"):
                    for child in tree.tail:
                        self.visit(child, level+1)
            else:
                if not tree.head.startswith("implicit_autogenerated_"):
                    for child in tree.tail:
                        self.visit(child, level)
        else:
            self.result += '\n' + '\t' * level + '(' + repr(tree) + ')'

    def dump(self):
        return self.result