File: nx_agraph.py

package info (click to toggle)
python-networkx 1.7~rc1-3
  • links: PTS, VCS
  • area: main
  • in suites: wheezy
  • size: 4,128 kB
  • sloc: python: 44,557; makefile: 135
file content (444 lines) | stat: -rw-r--r-- 13,121 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
"""
***************
Graphviz AGraph
***************

Interface to pygraphviz AGraph class.

Examples
--------
>>> G=nx.complete_graph(5)
>>> A=nx.to_agraph(G)
>>> H=nx.from_agraph(A)

See Also
--------
Pygraphviz: http://networkx.lanl.gov/pygraphviz


"""
__author__ = """Aric Hagberg (hagberg@lanl.gov)"""
#    Copyright (C) 2004-2008 by
#    Aric Hagberg <hagberg@lanl.gov>
#    Dan Schult <dschult@colgate.edu>
#    Pieter Swart <swart@lanl.gov>
#    All rights reserved.
#    BSD license.

__all__ = ['from_agraph', 'to_agraph',
           'write_dot', 'read_dot',
           'graphviz_layout',
           'pygraphviz_layout',
           'view_pygraphviz']

import os
import sys
import tempfile

import networkx as nx


def from_agraph(A,create_using=None):
    """Return a NetworkX Graph or DiGraph from a PyGraphviz graph.

    Parameters
    ----------
    A : PyGraphviz AGraph
      A graph created with PyGraphviz

    create_using : NetworkX graph class instance
      The output is created using the given graph class instance

    Examples
    --------
    >>> K5=nx.complete_graph(5)
    >>> A=nx.to_agraph(K5)
    >>> G=nx.from_agraph(A)
    >>> G=nx.from_agraph(A)


    Notes
    -----
    The Graph G will have a dictionary G.graph_attr containing
    the default graphviz attributes for graphs, nodes and edges.

    Default node attributes will be in the dictionary G.node_attr
    which is keyed by node.

    Edge attributes will be returned as edge data in G.  With
    edge_attr=False the edge data will be the Graphviz edge weight
    attribute or the value 1 if no edge weight attribute is found.

    """
    if create_using is None:
        if A.is_directed():
            if A.is_strict():
                create_using=nx.DiGraph()
            else:
                create_using=nx.MultiDiGraph()
        else:
            if A.is_strict():
                create_using=nx.Graph()
            else:
                create_using=nx.MultiGraph()

    # assign defaults
    N=nx.empty_graph(0,create_using)
    N.name=str(A)
    # add nodes, attributes to N.node_attr
    for n in A.nodes():
        str_attr=dict((str(k),v) for k,v in n.attr.items())
        N.add_node(str(n),**str_attr)

    # add edges, assign edge data as dictionary of attributes
    for e in A.edges():
        u,v=str(e[0]),str(e[1])
        attr=dict(e.attr)
        str_attr=dict((str(k),v) for k,v in attr.items())
        if not N.is_multigraph():
            if e.name is not None:
                str_attr['key']=e.name
            N.add_edge(u,v,**str_attr)
        else:
            N.add_edge(u,v,key=e.name,**str_attr)

    # add default attributes for graph, nodes, and edges
    # hang them on N.graph_attr
    N.graph['graph']=dict(A.graph_attr)
    N.graph['node']=dict(A.node_attr)
    N.graph['edge']=dict(A.edge_attr)
    return N

def to_agraph(N):
    """Return a pygraphviz graph from a NetworkX graph N.

    Parameters
    ----------
    N : NetworkX graph
      A graph created with NetworkX

    Examples
    --------
    >>> K5=nx.complete_graph(5)
    >>> A=nx.to_agraph(K5)

    Notes
    -----
    If N has an dict N.graph_attr an attempt will be made first
    to copy properties attached to the graph (see from_agraph)
    and then updated with the calling arguments if any.

    """
    try:
        import pygraphviz
    except ImportError:
        raise ImportError("to_agraph() requires pygraphviz: http://nx.lanl.gov/pygraphviz (not available for Python3")
    directed=N.is_directed()
    strict=N.number_of_selfloops()==0 and not N.is_multigraph()
    A=pygraphviz.AGraph(name=N.name,strict=strict,directed=directed)

    # default graph attributes
    A.graph_attr.update(N.graph.get('graph',{}))
    A.node_attr.update(N.graph.get('node',{}))
    A.edge_attr.update(N.graph.get('edge',{}))

    # add nodes
    for n,nodedata in N.nodes(data=True):
        A.add_node(n,**nodedata)

    # loop over edges

    if N.is_multigraph():
        for u,v,key,edgedata in N.edges_iter(data=True,keys=True):
            str_edgedata=dict((k,str(v)) for k,v in edgedata.items())
            A.add_edge(u,v,key=str(key),**str_edgedata)
    else:
        for u,v,edgedata in N.edges_iter(data=True):
            str_edgedata=dict((k,str(v)) for k,v in edgedata.items())
            A.add_edge(u,v,**str_edgedata)


    return A

def write_dot(G,path):
    """Write NetworkX graph G to Graphviz dot format on path.

    Parameters
    ----------
    G : graph
       A networkx graph
    path : filename
       Filename or file handle to write.

    """
    try:
        import pygraphviz
    except ImportError:
        raise ImportError("write_dot() requires pygraphviz: http://networkx.lanl.gov/pygraphviz (not available for Python3")

    A=to_agraph(G)
    A.write(path)
    A.clear()
    return

def read_dot(path):
    """Return a NetworkX graph from a dot file on path.

    Parameters
    ----------
    path : file or string
       File name or file handle to read.
    """
    try:
        import pygraphviz
    except ImportError:
        raise ImportError("read_dot() requires pygraphviz: http://networkx.lanl.gov/pygraphviz (not available for Python3)")
    A=pygraphviz.AGraph(file=path)
    return from_agraph(A)


def graphviz_layout(G,prog='neato',root=None, args=''):
    """Create node positions for G using Graphviz.

    Parameters
    ----------
    G : NetworkX graph
      A graph created with NetworkX
    prog : string
      Name of Graphviz layout program
    root : string, optional
      Root node for twopi layout
    args : string, optional
      Extra arguments to Graphviz layout program

    Returns : dictionary
      Dictionary of x,y, positions keyed by node.

    Examples
    --------
    >>> G=nx.petersen_graph()
    >>> pos=nx.graphviz_layout(G)
    >>> pos=nx.graphviz_layout(G,prog='dot')

    Notes
    -----
    This is a wrapper for pygraphviz_layout.

    """
    return pygraphviz_layout(G,prog=prog,root=root,args=args)

def pygraphviz_layout(G,prog='neato',root=None, args=''):
    """Create node positions for G using Graphviz.

    Parameters
    ----------
    G : NetworkX graph
      A graph created with NetworkX
    prog : string
      Name of Graphviz layout program
    root : string, optional
      Root node for twopi layout
    args : string, optional
      Extra arguments to Graphviz layout program

    Returns : dictionary
      Dictionary of x,y, positions keyed by node.

    Examples
    --------
    >>> G=nx.petersen_graph()
    >>> pos=nx.graphviz_layout(G)
    >>> pos=nx.graphviz_layout(G,prog='dot')

    """
    try:
        import pygraphviz
    except ImportError:
        raise ImportError("pygraphviz_layout() requires pygraphviz: http://networkx.lanl.gov/pygraphviz (not available for Python3")
    A=to_agraph(G)
    if root is not None:
        args+="-Groot=%s"%root
    A.layout(prog=prog,args=args)
    node_pos={}
    for n in G:
        node=pygraphviz.Node(A,n)
        try:
            xx,yy=node.attr["pos"].split(',')
            node_pos[n]=(float(xx),float(yy))
        except:
            print("no position for node",n)
            node_pos[n]=(0.0,0.0)
    return node_pos

@nx.utils.open_file(5, 'w')
def view_pygraphviz(G, edgelabel=None, prog='dot', args='',
                       suffix='', path=None):
    """Views the graph G using the specified layout algorithm.

    Parameters
    ----------
    G : NetworkX graph
        The machine to draw.
    edgelabel : str, callable, None
        If a string, then it specifes the edge attribute to be displayed
        on the edge labels. If a callable, then it is called for each
        edge and it should return the string to be displayed on the edges.
        The function signature of `edgelabel` should be edgelabel(data),
        where `data` is the edge attribute dictionary.
    prog : string
        Name of Graphviz layout program.
    args : str
        Additional arguments to pass to the Graphviz layout program.
    suffix : str
        If `filename` is None, we save to a temporary file.  The value of
        `suffix` will appear at the tail end of the temporary filename.
    path : str, None
        The filename used to save the image.  If None, save to a temporary
        file.  File formats are the same as those from pygraphviz.agraph.draw.

    Returns
    -------
    path : str
        The filename of the generated image.
    A : PyGraphviz graph
        The PyGraphviz graph instance used to generate the image.

    Notes
    -----
    If this function is called in succession too quickly, sometimes the
    image is not displayed. So you might consider time.sleep(.5) between
    calls if you experience problems.

    """
    if not len(G):
        raise nx.NetworkXException("An empty graph cannot be drawn.")

    import pygraphviz

    # If we are providing default values for graphviz, these must be set
    # before any nodes or edges are added to the PyGraphviz graph object.
    # The reason for this is that default values only affect incoming objects.
    # If you change the default values after the objects have been added,
    # then they inherit no value and are set only if explicitly set.

    # to_agraph() uses these values.
    attrs = ['edge', 'node', 'graph']
    for attr in attrs:
        if attr not in G.graph:
            G.graph[attr] = {}

    # These are the default values.
    edge_attrs = {'fontsize': '10'}
    node_attrs = {'style': 'filled',
                  'fillcolor': '#0000FF40',
                  'height': '0.75',
                  'width': '0.75',
                  'shape': 'circle'}
    graph_attrs = {}

    def update_attrs(which, attrs):
        # Update graph attributes. Return list of those which were added.
        added = []
        for k,v in attrs.items():
            if k not in G.graph[which]:
                G.graph[which][k] = v
                added.append(k)

    def clean_attrs(which, added):
        # Remove added attributes
        for attr in added:
            del G.graph[which][attr]
        if not G.graph[which]:
            del G.graph[which]

    # Update all default values
    update_attrs('edge', edge_attrs)
    update_attrs('node', node_attrs)
    update_attrs('graph', graph_attrs)

    # Convert to agraph, so we inherit default values
    A = to_agraph(G)

    # Remove the default values we added to the original graph.
    clean_attrs('edge', edge_attrs)
    clean_attrs('node', node_attrs)
    clean_attrs('graph', graph_attrs)

    # If the user passed in an edgelabel, we update the labels for all edges.
    if edgelabel is not None:
        if not hasattr(edgelabel, '__call__'):
            def func(data):
                return ''.join(["  ", str(data[edgelabel]), "  "])
        else:
            func = edgelabel

        # update all the edge labels
        if G.is_multigraph():
            for u,v,key,data in G.edges_iter(keys=True, data=True):
                # PyGraphviz doesn't convert the key to a string. See #339
                edge = A.get_edge(u,v,str(key))
                edge.attr['label'] = str(func(data))
        else:
            for u,v,data in G.edges_iter(data=True):
                edge = A.get_edge(u,v)
                edge.attr['label'] = str(func(data))

    if path is None:
        ext = 'png'
        if suffix:
            suffix = '_%s.%s' % (suffix, ext)
        else:
            suffix = '.%s' % (ext,)
        path = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
    else:
        # Assume the decorator worked and it is a file-object.
        pass

    display_pygraphviz(A, path=path, prog=prog, args=args)

    return path.name, A

def display_pygraphviz(graph, path, format=None, prog=None, args=''):
    """Internal function to display a graph in OS dependent manner.

    Parameters
    ----------
    graph : PyGraphviz graph
        A PyGraphviz AGraph instance.
    path :  file object
        An already opened file object that will be closed.
    format : str, None
        An attempt is made to guess the output format based on the extension
        of the filename. If that fails, the value of `format` is used.
    prog : string
        Name of Graphviz layout program.
    args : str
        Additional arguments to pass to the Graphviz layout program.

    Notes
    -----
    If this function is called in succession too quickly, sometimes the
    image is not displayed. So you might consider time.sleep(.5) between
    calls if you experience problems.

    """
    if format is None:
        filename = path.name
        format = os.path.splitext(filename)[1].lower()[1:]
    if not format:
        # Let the draw() function use its default
        format = None

    # Save to a file and display in the default viewer.
    # We must close the file before viewing it.
    graph.draw(path, format, prog, args)
    path.close()
    nx.utils.default_opener(filename)

# fixture for nose tests
def setup_module(module):
    from nose import SkipTest
    try:
        import pygraphviz
    except:
        raise SkipTest("pygraphviz not available")