modelverse/interface/HUTN/test/graph_compilation_action_language/postproc.py

import json


class Counter(object):
    def __init__(self):
        self.ctr = 0

    def inc(self):
        self.ctr += 1

    def get(self):
        return self.ctr


def postproc(input):
    lines = input.split('\n')

    graph = dict()
    kinds = dict()

    # read the graph
    for l in lines:
        if not l:
            continue
        if l[0] == "N":
            # N [A-z_0-9]+
            rest = l[2:]
            name = rest
            graph[name] = []
        elif l[0] == "V":
            # V [A-z_0-9]+\(.*?\)
            rest = l[2:]
            b = rest.find('(')
            name = rest[:b]
            kind = rest[b+1:-1]
            kinds[name] = kind
            graph[name] = []
        elif l[0] == "D":
            # D [A-z_0-9]+,".*?",[A-z_0-9]+
            rest = l[2:]
            c1 = rest.find(',')
            c2 = rest.find(',', c1+1)
            origin = rest[:c1]
            label = rest[c1+1:c2]
            target = rest[c2+1:]
            graph[origin].append((label, target))

    # order edges by label
    for n, edges in graph.iteritems():
        edges.sort(key=lambda x:x[0])

    renames = dict()
    new_graph = dict()
    new_kinds = dict()
    visited = dict()
    free_id = Counter()


    def dfs(name):
        if visited.get(name, False) == True:
            return
        new_name = 'auto_%s' % (free_id.get())
        free_id.inc()
        renames[name] = new_name
        new_graph[new_name] = []
        kind = kinds.get(name, None)
        if kind:
            new_kinds[new_name] = kind
        visited[name] = True
        edges = graph[name]
        for label, target in edges:
            if target[0] != '?':
                dfs(target)

    # rename nodes
    dfs('auto_initial_IP')

    # rename edges
    for n, edges in graph.iteritems():
        new_name = renames[n]
        for label, target in edges:
            if target[0] != '?':
                new_graph[new_name].append((label, renames[target]))
            else:
                new_graph[new_name].append((label, target))

    # sort nodes by name
    new_graph = new_graph.items()
    new_graph.sort(key=lambda x:x[0])

    f = open('output.dot', 'w')
    f.write('digraph G {\n')
    for n, edges in new_graph:
        kind = new_kinds.get(n, None)
        if kind:
            l = "{0}({1})".format(n, kind)
            f.write('{}[label={}];\n'.format(n, json.dumps(l)))
        else:
            f.write('{0}[label="{0}"];\n'.format(n))

    # include_id = 0
    for n, edges in new_graph:
        for label, target in edges:
            if target[0] == "?":
                target = target.replace('?', '_').replace('/', '_')
                # include_id += 1
            # t = target.replace('?', '').replace('/', '') + str(include_id)
            # include_id += 1
            f.write('{} -> {} [label={}];\n'.format(n, target, json.dumps(label)))
    f.write('}\n')

    output = []

    # output new file
    for n, edges in new_graph:
        kind = new_kinds.get(n, None)
        if kind:
            output.append("V %s(%s)\n" % (n, kind))
        else:
            output.append("N %s\n" % (n))

    for n, edges in new_graph:
        for label, target in edges:
            output.append("D %s,%s,%s\n" % (n, label, target))

    output = ''.join(output)
    output = output.replace("auto_0", "auto_initial_IP")

    return output


if __name__ == "__main__":
    import sys
    input = open(sys.argv[1], 'r').read()
    output = postproc(input)
    print output