Pretty print Scala complex object structures printed by specs2

scala_pretty_obj.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

__author__ = 'Yoann Bentz'
__version__ = '1.0'
__email__ = 'yoann@yoone.eu'

import os
import sys
import tempfile
import re
from subprocess import call


class Node:
    def __init__(self):
        self.value = None
        self.children = []
        self.can_have_children = False

    def add_child(self, node):
        self.children.append(node)

    def __str__(self):
        print(repr(self))

    def __repr__(self):
        return '{}({})'.format(self.value, len(self.children))


class State:
    stack = []
    in_val = False

    def stack_push(self, node):
        if self.stack_empty():
            self.root = node
        self.stack.append(node)

    def stack_pop(self):
        return self.stack.pop()

    def stack_head(self):
        return self.stack[-1]

    def stack_empty(self):
        return not self.stack


def parse_obj(string):
    state = State()

    for c in string.strip('\r\n\t '):
        if c == '(':
            state.in_val = False
            state.node.can_have_children = True
            if not state.stack_empty():
                state.stack_head().add_child(state.node)
            state.stack_push(state.node)
        elif c == ')':
            state.in_val = False
            head = state.stack_pop()
            if state.node and not state.node.can_have_children:
                head.add_child(state.node)
                state.node = None
        elif c == ',':
            state.in_val = False
            if state.node and not state.node.can_have_children:
                state.stack_head().add_child(state.node)
        else:
            if state.in_val:
                state.node.value += c
            elif c == ' ':
                continue  # skipping spaces after commas
            else:
                state.in_val = True
                state.node = Node()
                state.node.value = c

    return state.root


def prettify(node, indent=''):
    out = indent + node.value
    if node.can_have_children:
        out += '('
        if len(node.children) > 0:
            out += '\n'
            for child in node.children:
                out += prettify(child, indent + '  ')
            out += indent + ')\n'
        else:
            out += ')\n'
    else:
        out += '\n'
    return out


def split_dirty_input(s):
    s = re.sub(r'[\n\r]+', '', s)
    s = re.sub(r'^\s*(\[error\])?\s*[a-zA-Z: ]* expected ', '', s)
    s = re.sub(r'\([a-zA-Z]+\.scala:[0-9]+\)\s*$', '', s)
    return s.split(' found ')  # actual: ', found' but assuming comma might not be present


if __name__ == '__main__':
    if (len(sys.argv) < 2):
        print('usage: {} INPUTFILE'.format(sys.argv[0]))
        sys.exit(1)

    files = []
    with open(sys.argv[1]) as f:
        for obj_str in split_dirty_input(f.read()):
            pretty = prettify(parse_obj(obj_str))

            tmp_file = tempfile.mktemp()
            with open(tmp_file, 'w') as f:
                f.write(pretty)

            files.append(tmp_file)

    cmd = ['vimdiff'] + files
    call(cmd)

    for f in files:
        os.remove(f)