build_call_tree.py

#!/usr/bin/python3
import pydot
import sys


def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)


'''
Removes characters unrelated to tree description generated by CCTree.
Also skips all lines without '+->' or '+-<', consider them comments.
'''
def cleanup_cctree_lines(lines):
    new_lines=[]
    for i, line in enumerate(lines):
        j = line.find('+-')
        if j == -1:
            continue

        # Remove accidental leading '!' and '#'
        if line[0] == '!' or line[0] == '#':
            line = line[1:]
            j=j-1

        # Remove accidental '@' before '+-<' and '+->'
        k = line[:j].find('@')
        if k != -1:
            line = line[:k] + line[k+1:]

        new_lines.append(line)

    return new_lines


def name_to_color(color):
    if color == 'orange':
        return '#ff9900'
    if color == 'blue':
        return '#66ccff'
    if color == 'green':
        return '#99cc00'
    if color == 'red':
        return '#ff6666'
    if color == 'gray':
        return '#c0c0c0'
    if color == 'pink':
        return '#ff99cc'
    if color == 'purple':
        return '#cc99ff'
    return color


'''
This is a node/edge labeling format:
  +-> node a
    +-> node b # comment for edge a->b
Or:
  +-< node a
    +-< node b # comment for edge b->a
Edge comment should be separated by '#'.

Adding shape=some_shape in comments will affect node shape.
'''
def data_split(data):
    comment = ''
    shape = None
    color = None

    if '#' in data:
        name, comment = data.split('#', 1)
        name = name.strip()
        comments = comment.split()
        for i in range(len(comments)):
            if comments[i].startswith('shape='):
                shape = comments[i][6:].strip()
                comments = comments[:i] + comments[i + 1:]
                break
            if comments[i].startswith('color='):
                color = name_to_color(comments[i][6:].strip())
                comments = comments[:i] + comments[i + 1:]
                break
        comment = ' '.join(comments)
    else:
        name = data.strip()

    return name, comment, shape, color


def cctree_to_link_nodes_and_edges(path):
    with open(path) as fle:
        lines = cleanup_cctree_lines(fle.readlines())

        if not len(lines):
            eprint('Empty CCTree file')
            return [], []

        j = lines[0].find('+-')
        if j == -1 or j + 2 >= len(lines[0]):
            eprint('Bad format of CCTree file')
            return [], []

        offs = []
        for line in lines:
            off = line.find('+-')
            if off == -1:
                eprint('Bad format of CCTree file:', line)
                return [], []
            offs.append(off)

        nodes = {}
        for i, _ in enumerate(lines):
            node, _, shape, color = data_split(lines[i][offs[i]+4:].strip())
            if node not in nodes:
                nodes[node] = {}
            if shape and 'shape' not in nodes[node]:
                nodes[node]['shape'] = shape
            if color and 'color' not in nodes[node]:
                nodes[node]['color'] = color

        links = []
        root = 0
        up = True
        for i, _ in enumerate(lines):
            found = False
            for j in range(i - 1, root-1, -1):
                if offs[j] >= offs[i]:
                    continue
                elif offs[j] + 2 != offs[i]:
                    eprint('Bad format of CCTree file:', ''.join(lines[j:i+1]))
                    return [], []
                else:
                    found = True
                    a, comment_a, _, _ = data_split(lines[i][offs[i]+4:].strip())

                    b, _, _, _ = data_split(lines[j][offs[j]+4:].strip())
                    edge = []
                    if up:
                        edge=[a, b, comment_a]
                    else:
                        edge=[b, a, comment_a]

                    if edge not in links:
                        links.append(edge)
                    break;

            '''
            Next root node. This happens when we get to next CCTree output
            '''
            if not found:
                root = i
                if lines[i][offs[i] + 2] == '<':
                    up = True
                elif lines[i][offs[i] + 2] == '>':
                    up = False
                else:
                    eprint('Bad format of CCTree file:', lines[0])
                    return [], []

        return list(nodes.items()), links


def nodes_and_edges_to_tree(nodes, links, dot_path):
    gr = pydot.Dot(graph_type='digraph')
    counter = 0
    nodemap = {}

    for node, options in nodes:
        shape = 'rectangle'
        if 'shape' in options:
            shape = options['shape']
        color = '"#ffcc00"'
        if 'color' in options:
            color = '"' + options['color'] + '"'

        if node in nodemap:
            continue
        nodemap[node] = str(counter)
        counter += 1

        nod = pydot.Node(nodemap[node], label=node, shape=shape, color=color)
        gr.add_node(nod)
    for link in links:
        edge = pydot.Edge(nodemap[link[0]], nodemap[link[1]], label=link[2])
        gr.add_edge(edge)
    gr.write(dot_path)


if __name__ == '__main__':
    if len(sys.argv) < 2 or len(sys.argv) > 3:
        eprint("usage: %s <path/to/graph.cct> [<path/to/graph.dot>]\n" % (sys.argv[0]))
        exit(1)

    dot_path = '/dev/stdout'
    if len(sys.argv) == 3:
        dot_path = sys.argv[2]

    nodes, links = cctree_to_link_nodes_and_edges(sys.argv[1])

    nodes_and_edges_to_tree(nodes, links, dot_path)