polygraph/polygraph.py

49 lines
1.4 KiB
Python
Executable File

#!/usr/bin/python
import pygraphviz as pgv
import json
import argparse
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('--color', '-c', action='store_true', help='Include node colors 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
for user in data['members']:
node_color = 'white'
if settings.color and user.has_key('color'):
node_color = user['color']
graph.add_node(user['name'], fillcolor=node_color)
# Add edges to graph
edge_style = 'solid'
for (member1,member2,edge) in data['relationships']:
edge_style = edge
graph.add_edge(member1, member2, style=edge_style, color='black')
graph.layout()
graph.draw(settings.output)
if __name__ == '__main__': main()