EQ graph builder

graphs.py

import json
import click
from graphviz import Digraph
from collections import OrderedDict

class Graph:
    def __init__(self, filename):
        with open(filename, 'r') as f:
            self.input = json.loads(f.read())

        #graph [label="Orthogonal edges", splines=ortho, nodesep=0.8]
        """
        graph_attrs = [
            ('label', 'Orthogonal edges'),
            ('splines', 'ortho'),
            ('nodesep', '0.8')
        ]
        """
        graph_attrs = []
        self.dot = Digraph(comment=filename, graph_attr=graph_attrs)

        self.blocks_by_id = self.load_blocks_by_id()

        self.groups = self.load_groups_by_id()

        self.block_ids = []

        for block_id in self.blocks_by_id:
            self.dot.node(block_id)
            self.block_ids.append(block_id)

        for block_id, block in self.blocks_by_id.items():
            routing_rules = block.get('routing_rules', [])

            if not routing_rules:
                try:
                    next_block_index = self.block_ids.index(block_id) + 1
                    next_block = self.block_ids[next_block_index]
                except IndexError:
                    print("Failed to find next block for: {}".format(block_id))
                self.dot.edge(block_id, next_block)
            for routing_rule in routing_rules:
                if routing_rule.get('goto'):

                    if routing_rule['goto'].get('block'):
                        target = routing_rule['goto']['block']
                    elif routing_rule['goto'].get('group'):
                        target = self.groups[routing_rule['goto']['group']]['blocks'][0]['id']
                    else:
                        print("Unknown routing rule: {}".format(routing_rule))
                        continue

                    when_rules = routing_rule['goto'].get('when')

                    if not when_rules:
                        rule = 'DEFAULT'
                        self.dot.edge(block['id'], target, label=rule)
                    else:
                        for when_rule in when_rules:
                            rule = ' '.join([str(when_rule.get('id', '')),
                                             str(when_rule.get('condition', '')),
                                             str(when_rule.get('value', ''))])
                            self.dot.edge(block['id'], target, label=rule)


    def load_groups_by_id(self):
        groups = OrderedDict()

        for section in self.input['sections']:
            for group in section['groups']:
                groups[group['id']] = group

        return groups


    def load_blocks_by_id(self):
        blocks = OrderedDict()

        for section in self.input['sections']:
            for group in section['groups']:
                for block in group['blocks']:
                    blocks[block['id']] = block

        return blocks


@click.command()
@click.argument('filename')
def build_graph(filename):
    graph = Graph(filename)
    graph.dot.render('/tmp/g.gv', view=True)
    print("Done!")


if __name__ == '__main__':
    build_graph()