Code to help transform a DOT file (graph)

graph_helpers.py

# This module is meant to be imported.

# You'll need to run:
#   pip install pydot

'''
# Example usage:

import graph_helpers as gh

gh.load_graph()
g = gh.get_graph()
gh.populate_missing_nodes()

# Add 'delete_me' attribute to edges and nodes as you see fit...
for n in g.get_nodes():
    if n.get_attributes().get('color') == 'red':
        n.get_attributes()['delete_me'] = 'true'
for e in g.get_edges():
    if e.get_attributes().get('color') == 'red':
        e.get_attributes()['delete_me'] = 'true'

gh.delete_flagged_nodes_and_edges() # Actually delete the edges and nodes that have the 'delete_me' attribute.
gh.delete_nodes_with_zero_edges()

print g.to_string()
'''

import sys

try:
    import pydot
except ImportError:
    sys.stderr.write("ImportError: pydot\n")
    sys.stderr.write("Please run:\n")
    sys.stderr.write("  pip install pydot\n")
    sys.stderr.flush()
    sys.exit(1)

g = None # Global graph object. Call load_graph() to populate it.

def load_graph(filepath):
    with open(filepath) as f:
        dot_text = f.read()

    graph_list = pydot.graph_from_dot_data(dot_text)
    assert len(graph_list) == 1

    global g
    g = graph_list[0]

def get_graph():
    return g

def populate_missing_nodes():
    '''
    Adds any nodes that are not explicitly stated.

    For example, this would turn the graph
        digraph G {
            a -> b;
        }
    into
        digraph G {
            a;
            b;
            a -> b;
        }
    '''
    for e in g.get_edges():
        for node_name in (e.get_source(), e.get_destination()):
            if not g.get_node(node_name):
                g.add_node(pydot.Node(name = node_name))

def delete_flagged_nodes_and_edges():
    '''
    Deletes any node or edge that has 'delete_me' attribute set (to anything).
    When deleting a node, it also deletes any edges connected to that node.
    '''
    for e in g.get_edges():
        if ('delete_me' in e.get_attributes()
                or 'delete_me' in get_source(e).get_attributes()
                or 'delete_me' in get_dest(e).get_attributes()
            ):
            g.del_edge(e.get_source(), e.get_destination())

    for n in g.get_nodes():
        if 'delete_me' in n.get_attributes():
            g.del_node(n)

def delete_nodes_with_zero_edges():
    '''
    Deletes any node that has no inbound or outbound edges.
    '''
    for n in g.get_nodes():
        if n.get_name() not in ('graph', 'node', 'edge') and len(get_all_edges(n)) == 0:
            g.del_node(n)


# --------------------
# Helper functions:
# --------------------

def get_node(node_name):
    node_list = g.get_node(node_name)
    assert len(node_list) in (0, 1)
    if node_list:
        return node_list[0]
    else:
        return pydot.Node(name = node_name)

def get_source(e):
    return get_node(e.get_source())

def get_dest(e):
    return get_node(e.get_destination())

def get_inbound_edges(node):
    ''' Returns a list of all edges whose destination is node. '''
    if isinstance(node, pydot.Node):
        node = node.get_name()
    return [
        e for e in g.get_edges()
        if e.get_destination() == node
    ]

def get_outbound_edges(node):
    ''' Returns a list of all edges whose source is node. '''
    if isinstance(node, pydot.Node):
        node = node.get_name()
    return [
        e for e in g.get_edges()
        if e.get_source() == node
    ]

def get_all_edges(node):
    return get_inbound_edges(node) + get_outbound_edges(node)

# --------------------
# End helper functions
# --------------------

render_graph.sh

#!/usr/bin/env bash

# This script opens a window to display a rendering of your DOT file.

# You'll need to install graphviz first:
#   brew install graphviz

set -e

if [ ! -f "$1" ]; then
    echo "File not found!"
    echo "Usage: ./render_graph.sh filename.dot [program_to_use]"
    exit
fi

program=${2:-dot}

file_name_only=`basename "$1"`
epoch_time=`date +%s`
random_number=`echo $RANDOM`
temp_file="/tmp/${file_name_only}.${epoch_time}.${random_number}.pdf"

$program "$1" -Tpdf -o "$temp_file"
open "$temp_file"