Getting started with Graphviz and Python

About the graphviz module

Installation: sudo pip install graphviz

Hello Graphviz

First we need to import the graphviz module:

import graphviz as gv

Now we can create the graph object g1 and add two nodes A and B as well as an edge to connect the two.

g1 = gv.Graph(format='png')
g1.node('A')
g1.node('B')
g1.edge('A', 'B')

Let's have a look at the the dot code that this will generate behind the scenes:

print(g1.source)
graph {
    A
    B
        A -- B
}

That looks about right. Let's render it:

filename = g1.render(filename='img/p0204_1')
print(filename)
img/p0204_1.png

As you can see, the graphviz module takes care of appending the correct file extension. Let's have a look at the result:

g1

Directed Graphs

To create a directed version of the previous graph we only need to replace the Graph object with a Digraph.

g2 = gv.Digraph(format='png')
g2.node('A')
g2.node('B')
g2.edge('A', 'B')
g2.render('img/p0204_2')

g2

A more concise approach

The above approach is fine for small graphs but for larger ones we should create a few helper functions. Since we will stick to PNG as our output format there is no need to explicity specifiy it for every graph.

import functools
graph = functools.partial(gv.Graph, format='png')
digraph = functools.partial(gv.Digraph, format='png')

Now creating a new graph is as easy as calling graph()

g3 = graph()

Next we would like to specify the nodes and edges as simple data structures instead typing loads of method calls.

nodes = ['A', 'B', ('C', {})]

Every node is represented by either a string or a tuple. In the latter case the second element of the tuple holds all attributes of that node as we will see later.

The represenation of edges is quite similar:

edges = [
    ('A', 'B'),
    ('B', 'C'),
    (('A', 'C'), {}),
]

To actually create nodes and edges from these data structures we implement the functions add_nodes and add_edges. Note that both functions return the updated graph which will come in handy in just a second.

def add_nodes(graph, nodes):
    for n in nodes:
        if isinstance(n, tuple):
            graph.node(n[0], **n[1])
        else:
            graph.node(n)
    return graph


def add_edges(graph, edges):
    for e in edges:
        if isinstance(e[0], tuple):
            graph.edge(*e[0], **e[1])
        else:
            graph.edge(*e)
    return graph

Now we can create graphs in a much more concise way:

add_edges(
    add_nodes(digraph(), ['A', 'B', 'C']),
    [('A', 'B'), ('A', 'C'), ('B', 'C')]
).render('img/p0204_4')

g4

Labels

Labeling nodes and edges is as easy as adding an entry to their dictionaries of attributes.

add_edges(
    add_nodes(digraph(), [
        ('A', {'label': 'Node A'}),
        ('B', {'label': 'Node B'}),
        'C'
    ]),
    [
        (('A', 'B'), {'label': 'Edge 1'}),
        (('A', 'C'), {'label': 'Edge 2'}),
        ('B', 'C')
    ]
).render('img/p0204_5')

g5

Styling

Graphiz let's you style graphs extensively by setting specific attributes. The following example will showcase a few of these attributes. For a complete list please refer to http://www.graphviz.org/doc/info/attrs.html.

First let's create an unstyled graph:

g6 = add_edges(
    add_nodes(digraph(), [
        ('A', {'label': 'Node A'}),
        ('B', {'label': 'Node B'}),
        'C'
    ]),
    [
        (('A', 'B'), {'label': 'Edge 1'}),
        (('A', 'C'), {'label': 'Edge 2'}),
        ('B', 'C')
    ]
)

We will use a simple dictionary to specifiy the attributes that we want to change:

styles = {
    'graph': {
        'label': 'A Fancy Graph',
        'fontsize': '16',
        'fontcolor': 'white',
        'bgcolor': '#333333',
        'rankdir': 'BT',
    },
    'nodes': {
        'fontname': 'Helvetica',
        'shape': 'hexagon',
        'fontcolor': 'white',
        'color': 'white',
        'style': 'filled',
        'fillcolor': '#006699',
    },
    'edges': {
        'style': 'dashed',
        'color': 'white',
        'arrowhead': 'open',
        'fontname': 'Courier',
        'fontsize': '12',
        'fontcolor': 'white',
    }
}

Now we need another helper function to actually apply those attributes to the graph:

def apply_styles(graph, styles):
    graph.graph_attr.update(
        ('graph' in styles and styles['graph']) or {}
    )
    graph.node_attr.update(
        ('nodes' in styles and styles['nodes']) or {}
    )
    graph.edge_attr.update(
        ('edges' in styles and styles['edges']) or {}
    )
    return graph

Let's try it out!

g6 = apply_styles(g6, styles)
g6.render('img/p0204_6')

g6

Subgraphs

Finally, we will have a look a subgraphs. First, let's create two seperate graphs g7 and g8, each with their own styles.

g7 = add_edges(
    add_nodes(digraph(), [
        ('A', {'label': 'Node A'}),
        ('B', {'label': 'Node B'}),
        'C'
    ]),
    [
        (('A', 'B'), {'label': 'Edge 1'}),
        (('A', 'C'), {'label': 'Edge 2'}),
        ('B', 'C')
    ]
)

g8 = apply_styles(
    add_edges(
        add_nodes(digraph(), [
            ('D', {'label': 'Node D'}),
            ('E', {'label': 'Node E'}),
            'F'
        ]),
        [
            (('D', 'E'), {'label': 'Edge 3'}),
            (('D', 'F'), {'label': 'Edge 4'}),
            ('E', 'F')
        ]
    ),
    {
        'nodes': {
            'shape': 'square',
            'style': 'filled',
            'fillcolor': '#cccccc',
        }
    }
)

Now we combine the two graphs by making g8 a subgraph of g7 and adding an extra edge to connect the two:

g7.subgraph(g8)
g7.edge('B', 'E', color='red', weight='2')

Let's have a look at the result:

g7.render('img/p0204_7')

g7

Learn more

About this Entry

Related Content


Comments