nsq_graph.py

import json
import argparse

import pygraphviz as pgv


def mk_key(topic, channel, channels, host_class, other_host_class):
    if channel == 'nsq_to_file' and len(channels) > 1:
        # don't show archival of topics already covered
        return None
    if channel == 'nsq_to_file':
        return "nsq_to_file (%s)" % topic
    return channel


def mk_pgv_graph(data, key_fxn=mk_key):  # noqa
    """
    Takes nsq graph data (as generated by nsq_data.py) and makes a pygraphviz graph.

    `key_fxn` takes info about a potential edge (see `mk_key` for args) and returns
    a "key" to label that edge. If it returns `None`, we skip drawing the edge
    """
    G = pgv.AGraph(strict=False, directed=True)
    for host_class in data:
        G.add_node(host_class)
    for host_class, topics in data.items():
        for topic, channels in topics.items():
            for channel, other_host_classes in channels.items():
                for other in other_host_classes:
                    key = key_fxn(topic, channel, channels, host_class, other)
                    if key is None:
                        continue
                    try:
                        G.get_edge(host_class, other, key=key)
                    except KeyError:
                        # edge not present; add it
                        G.add_edge(host_class, other, key=key, label=key)
    return G


def draw_pgv_graph(G, output_path):
    G.layout(prog='dot')
    G.draw(output_path)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=argparse.FileType('r'))
    parser.add_argument('--output', required=True)
    args = parser.parse_args()
    data = json.load(args.input)
    dg = mk_pgv_graph(data)
    draw_pgv_graph(dg, args.output)