polygraph/polygraph.py

#!/usr/bin/python

import pygraphviz as pgv
import json
import argparse

# These are the attributes to use for color-coding the graphs

# Three gender colours, so that we encapsulate masculine, feminine, and a catch-all for other identities
# If you'd rather explicitly colour-code other identities, just add them :)
gender_colors = {
    'm': 'blue',
    'f': 'pink',
    'o': 'green',
    }


def parse_args():
    parser = argparse.ArgumentParser(description='Generate graphs of polyamorous family networks from data files')
    parser.add_argument('--input', '-i', default='network.json', help='Input file in json format')
    parser.add_argument('--output', '-o', default='network.png', help='Output image file. Format auto-detected from file extension')
    parser.add_argument('--gender', '-g', action='store_true', help='Indicate gender in output graph')
    parser.add_argument('--relationships', '-r', action='store_true', help='Indicate relationship details (std risk, relationship type) in output graph')
    return parser.parse_args()


def main():
    settings = parse_args()

    try:
        data = json.load(open(settings.input, 'r'))
    except Exception as e:
        print "Error parsing data file: " + str(e)
        return

    graph = pgv.AGraph()
    graph.graph_attr['overlap'] = 'prism'
    graph.node_attr['style'] = 'filled'
    graph.edge_attr['arrowsize'] = 0.5

    # Add nodes to graph
    gender_color = 'white'

    for user in data['members']:
        if settings.gender:
            gender_color = gender_colors[user['gender']]

        graph.add_node(user['name'], fillcolor=gender_color)

    # Add edges to graph
    edge_style = 'solid'

    for (member1,member2,rel,edge) in data['relationships']:
        if settings.relationships:
            edge_style = edge
            relationship_color = relationship_colors[rel]

        graph.add_edge(member1, member2, style=edge_style, color='black')

    graph.layout()
    graph.draw(settings.output)
    

if __name__ == '__main__': main()