nx_pydot.py

"""
Provide functions for converting to/from NetworkX and pydot+graphviz.

Note: Much of this module exists in NetworkX 1.8.  NetworkX will probably
continue to support basic conversions between NetworkX and pydot+graphviz,
but visualization functionality may be discontinued in future versions.

pydot only supports Python versions 2.6+ but not the 3.x series.  A version
which supports 2.6+ and 3.x is available here:

   https://github.com/nlhepler/pydot

We use this version for a number of reasons:

    1) It supports 2.6+ and 3.x
    2) It has improved write support.
    3) It is based off pydot 1.0.28 (Jan 2012) which fixes:
         https://code.google.com/p/pydot/issues/detail?id=61

To simplify installation of nxpd, an updated version of pydot is included
and used internally.

See Also
--------
pydot: http://code.google.com/p/pydot/
Graphviz: http://www.research.att.com/sw/tools/graphviz/
DOT Language: http://www.graphviz.org/doc/info/lang.html

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

from __future__ import absolute_import

import os
import sys
import tempfile
import time

from . import pydot
from .params import nxpdParams
from .utils import get_fobj, default_opener, make_str

import networkx as nx

__author__ = "\n".join(["Aric Hagberg (aric.hagberg@gmail.com)",
                        "chebee7i (chebee7i@gmail.com)"])

__all__ = ['write_dot', 'read_dot', 'graphviz_layout', 'pydot_layout',
           'to_pydot', 'from_pydot', 'draw_pydot']

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

    Path can be a string or a file handle.

    """
    P = to_pydot(G)

    fobj, close = get_fobj(path, mode)
    try:
        fobj.write(P.to_string())
    finally:
        if close:
            fobj.close()

def read_dot(path, mode='r'):
    """Return a NetworkX MultiGraph or MultiDiGraph from a dot file on path.

    Parameters
    ----------
    path : filename or file handle

    Returns
    -------
    G : NetworkX multigraph
        A MultiGraph or MultiDiGraph.

    Notes
    -----
    Use G=nx.Graph(nx.read_dot(path)) to return a Graph instead of a MultiGraph.

    """
    fobj, close = get_fobj(path, mode)

    try:
        data = fobj.read()
    finally:
        if close:
            fobj.close()

    P = pydot.graph_from_dot_data(data)
    G = from_pydot(P)

    return G

def from_pydot(P):
    """Return a NetworkX graph from a Pydot graph.

    Parameters
    ----------
    P : Pydot graph
      A graph created with Pydot

    Returns
    -------
    G : NetworkX multigraph
        A MultiGraph or MultiDiGraph.

    Examples
    --------
    >>> K5 = nx.complete_graph(5)
    >>> A = nx.to_pydot(K5)
    >>> G = nx.from_pydot(A) # return MultiGraph
    >>> G = nx.Graph(nx.from_pydot(A)) # make a Graph instead of MultiGraph

    """
    if P.get_strict(None): # pydot bug: get_strict() shouldn't take argument
        multiedges = False
    else:
        multiedges = True

    if P.get_type() == 'graph': # undirected
        if multiedges:
            create_using = nx.MultiGraph()
        else:
            create_using = nx.Graph()
    else:
        if multiedges:
            create_using = nx.MultiDiGraph()
        else:
            create_using = nx.DiGraph()

    # assign defaults
    N = nx.empty_graph(0, create_using)
    N.name = P.get_name()

    # add nodes, attributes to N.node_attr
    for p in P.get_node_list():
        n = p.get_name().strip('"')
        if n in ('node', 'graph', 'edge'):
            continue
        N.add_node(n, **p.get_attributes())

    # add edges
    for e in P.get_edge_list():
        u = e.get_source()
        v = e.get_destination()
        attr = e.get_attributes()
        s = []
        d = []

        if isinstance(u, basestring):
            s.append(u.strip('"'))
        else:
            for unodes in u['nodes'].iterkeys():
                s.append(unodes.strip('"'))

        if isinstance(v, basestring):
            d.append(v.strip('"'))
        else:
            for vnodes in v['nodes'].iterkeys():
                d.append(vnodes.strip('"'))

        for source_node in s:
            for destination_node in d:
                N.add_edge(source_node, destination_node, **attr)

    # add default attributes for graph, nodes, edges
    N.graph['graph'] = P.get_attributes()
    try:
        N.graph['node'] = P.get_node_defaults()[0]
    except:# IndexError,TypeError:
        N.graph['node'] = {}
    try:
        N.graph['edge'] = P.get_edge_defaults()[0]
    except:# IndexError,TypeError:
        N.graph['edge'] = {}
    return N

def filter_attrs(attrs, attr_type):
    """
    Helper function to keep only pydot supported attributes.

    All unsupported attributes are filtered out.

    Parameters
    ----------
    attrs : dict
        A dictionary of attributes.
    attr_type : str
        The type of attributes. Must be 'edge', 'graph', or 'node'.

    Returns
    -------
    d : dict
        The filtered attributes.

    """
    if attr_type == 'edge':
        accepted = pydot.EDGE_ATTRIBUTES
    elif attr_type == 'graph':
        accepted = pydot.GRAPH_ATTRIBUTES
    elif attr_type == 'node':
        accepted = pydot.NODE_ATTRIBUTES
    else:
        raise Exception("Invalid attr_type.")

    d = dict( [(k,v) for (k,v) in attrs.items() if k in accepted] )
    return d

def to_pydot(G, raise_exceptions=True):
    """Return a pydot graph from a NetworkX graph G.

    All node names are converted to strings.  However, no preprocessing is
    performed on the edge/graph/node attribute values since some attributes
    need to be strings while other need to be floats. If pydot does not handle
    needed conversions, then your graph should be modified beforehand.

    Generally, the rule is:  If the attribute is a supported Graphviz
    attribute, then it will be added to the Pydot graph (and thus, assumed to
    be in the proper format for Graphviz).

    Parameters
    ----------
    G : NetworkX graph
        A graph created with NetworkX.
    raise_exceptions : bool
        If `True`, raise any exceptions.  Otherwise, the exception is ignored
        and the procedure continues.

    Examples
    --------
    >>> G = nx.complete_graph(5)
    >>> G.add_edge(2, 10, color='red')
    >>> P = nx.to_pydot(G)

    """
    # Set Graphviz graph type.
    if G.is_directed():
        graph_type = 'digraph'
    else:
        graph_type = 'graph'

    strict = G.number_of_selfloops() == 0 and not G.is_multigraph()

    # Create the Pydot graph.
    name = G.graph.get('name')
    graph_defaults = filter_attrs(G.graph, 'graph')
    if name is None:
        P = pydot.Dot(graph_type=graph_type, strict=strict, **graph_defaults)
    else:
        P = pydot.Dot(name, graph_type=graph_type, strict=strict,
                      **graph_defaults)

    # Set default node attributes, if possible.
    node_defaults = filter_attrs(G.graph.get('node', {}), 'node')
    if node_defaults:
        try:
            P.set_node_defaults(**node_defaults)
        except:
            if raise_exceptions:
                raise

    # Set default edge attributes, if possible.
    edge_defaults = filter_attrs(G.graph.get('edge', {}), 'edge')
    if edge_defaults:
        # This adds a node called "edge" to the graph.
        try:
            P.set_edge_defaults(**edge_defaults)
        except:
            if raise_exceptions:
                raise

    # Add the nodes.
    for n,nodedata in G.nodes(data=True):
        attrs = filter_attrs(nodedata, 'node')
        node = pydot.Node(make_str(n), **attrs)
        P.add_node(node)

    # Add the edges.
    if G.is_multigraph():
        for u,v,key,edgedata in G.edges(data=True,keys=True):
            attrs = filter_attrs(edgedata, 'edge')
            uu, vv, kk = make_str(u), make_str(v), make_str(key)
            edge = pydot.Edge(uu, vv, key=kk, **attrs)
            P.add_edge(edge)
    else:
        for u,v,edgedata in G.edges(data=True):
            attrs = filter_attrs(edgedata, 'edge')
            uu, vv = make_str(u), make_str(v)
            edge = pydot.Edge(uu, vv, **attrs)
            P.add_edge(edge)
    return P

def graphviz_layout(G, prog='neato', root=None, **kwds):
    """Create node positions using Pydot and Graphviz.

    Returns a dictionary of positions keyed by node.

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

    Notes
    -----
    This is a wrapper for pydot_layout.
    """
    return pydot_layout(G=G, prog=prog, root=root, **kwds)


def pydot_layout(G, prog='neato', root=None, **kwds):
    """Create node positions using Pydot and Graphviz.

    Returns a dictionary of positions keyed by node.

    Examples
    --------
    >>> G = nx.complete_graph(4)
    >>> pos = nx.pydot_layout(G)
    >>> pos = nx.pydot_layout(G, prog='dot')

    """
    P = to_pydot(G)
    if root is not None :
        P.set("root",make_str(root))

    # Need to verify this, but pydot encodes as utf-8.
    D = P.create_dot(prog=prog)

    if D == "":  # no data returned
        raise Exception("Graphviz layout with {0} failed.".format(prog))

    # And pydot decodes from utf-8. So the nodes in Q are unicode.
    Q = pydot.graph_from_dot_data(D)

    node_pos = {}
    for n in G.nodes():
        # pydot handles all encodings/decodings, so we keep it unicode.
        pydot_node = pydot.Node(make_str(n)).get_name()
        node = Q.get_node(pydot_node)

        if isinstance(node,list):
            node = node[0]
        pos = node.get_pos()[1:-1] # strip leading and trailing double quotes
        if pos != None:
            xx,yy = pos.split(",")
            node_pos[n] = (float(xx), float(yy))
    return node_pos

def draw_pydot(G, filename=None, format=None, prefix=None, suffix=None,
                  layout='dot', args=None, show=None):
    """Draws the graph G using pydot and graphviz.

    Parameters
    ----------
    G : graph
        A NetworkX graph object (e.g., Graph, DiGraph).

    filename : str, None, file object
        The name of the file to save the image to.  If None, save to a
        temporary file with the name:
             nx_PREFIX_RANDOMSTRING_SUFFIX.ext.
        File formats are inferred from the extension of the filename, when
        provided.  If the `format` parameter is not `None`, it overwrites any
        inferred value for the extension.

    format : str
        An output format. Note that not all may be available on every system
        depending on how Graphviz was built. If no filename is provided and
        no format is specified, then a 'png' image is created. Other values
        for `format` are:

            'canon', 'cmap', 'cmapx', 'cmapx_np', 'dia', 'dot',
            'fig', 'gd', 'gd2', 'gif', 'hpgl', 'imap', 'imap_np',
            'ismap', 'jpe', 'jpeg', 'jpg', 'mif', 'mp', 'pcl', 'pdf',
            'pic', 'plain', 'plain-ext', 'png', 'ps', 'ps2', 'svg',
            'svgz', 'vml', 'vmlz', 'vrml', 'vtx', 'wbmp', 'xdot', 'xlib'

    prefix : str | None
        If `filename` is None, we save to a temporary file.  The value of
        `prefix` will appear after 'nx_' but before random string
        and file extension. If None, then the graph name will be used.

    suffix : str | None
        If `filename` is None, we save to a temporary file.  The value of
        `suffix` will appear at after the prefix and random string but before
        the file extension. If None, then no suffix is used.

    layout : str
        The graphviz layout program.  Pydot is responsible for locating the
        binary. Common values for the layout program are:
            'neato','dot','twopi','circo','fdp','nop', 'wc','acyclic','gvpr',
            'gvcolor','ccomps','sccmap','tred'

    args : list
        Additional arguments to pass to the Graphviz layout program.
        This should be a list of strings.  For example, ['-s10', '-maxiter=10'].

    show : bool
        If `True`, then the image is displayed using the OS's default viewer
        after drawing it. If show equals 'ipynb', then the image is displayed
        inline for an IPython notebook.  If `None`, then the value of
        nxpdParams['show'] is used.  By default, it is set to `True`.

    """
    # Determine the output format
    if format is None:
        # grab extension from filename
        if filename is None:
            # default to png
            ext = 'png'
        else:
            ext = os.path.splitext(filename)[-1].lower()[1:]
    else:
        ext = format

    # Determine the "path" to be passed to pydot.Dot.write()
    if filename is None:
        if prefix is None:
            prefix = G.graph.get("name", '')

        if prefix:
            fn_prefix = "nx_{0}_".format(prefix)
        else:
            fn_prefix = "nx_"

        if suffix:
            fn_suffix = '_{0}.{1}'.format(suffix, ext)
        else:
            fn_suffix = '.{0}'.format(ext)

        fobj = tempfile.NamedTemporaryFile(prefix=fn_prefix,
                                           suffix=fn_suffix,
                                           delete=False)
        fname = fobj.name
        close = True
    else:
        fobj, close = get_fobj(filename, 'w+b')
        fname = fobj.name

    # Include additional command line arguments to the layout program.
    if args is None:
        args = []
        prog = layout
    else:
        args = list(args)
        prog = [layout] + args

    # Draw the image.
    G2 = to_pydot(G)
    G2.write(fobj, prog=prog, format=ext)
    if close:
        fobj.close()

    if show is None:
        show = nxpdParams['show']

    if show:
        if show == 'ipynb':
            from IPython.core.display import Image
            return Image(filename=fname, embed=True)
        else:
            default_opener(fname)
            if sys.platform == 'linux2':
                # necessary when opening many images in a row
                time.sleep(.5)

    return fname