1"""
2*****
3Pydot
4*****
5
6Import and export NetworkX graphs in Graphviz dot format using pydot.
7
8Either this module or nx_agraph can be used to interface with graphviz.
9
10Examples
11--------
12>>> G = nx.complete_graph(5)
13>>> PG = nx.nx_pydot.to_pydot(G)
14>>> H = nx.nx_pydot.from_pydot(PG)
15
16See Also
17--------
18 - pydot:         https://github.com/erocarrera/pydot
19 - Graphviz:      https://www.graphviz.org
20 - DOT Language:  http://www.graphviz.org/doc/info/lang.html
21"""
22from locale import getpreferredencoding
23from networkx.utils import open_file
24import networkx as nx
25
26__all__ = [
27    "write_dot",
28    "read_dot",
29    "graphviz_layout",
30    "pydot_layout",
31    "to_pydot",
32    "from_pydot",
33]
34
35
36@open_file(1, mode="w")
37def write_dot(G, path):
38    """Write NetworkX graph G to Graphviz dot format on path.
39
40    Path can be a string or a file handle.
41    """
42    P = to_pydot(G)
43    path.write(P.to_string())
44    return
45
46
47@open_file(0, mode="r")
48def read_dot(path):
49    """Returns a NetworkX :class:`MultiGraph` or :class:`MultiDiGraph` from the
50    dot file with the passed path.
51
52    If this file contains multiple graphs, only the first such graph is
53    returned. All graphs _except_ the first are silently ignored.
54
55    Parameters
56    ----------
57    path : str or file
58        Filename or file handle.
59
60    Returns
61    -------
62    G : MultiGraph or MultiDiGraph
63        A :class:`MultiGraph` or :class:`MultiDiGraph`.
64
65    Notes
66    -----
67    Use `G = nx.Graph(read_dot(path))` to return a :class:`Graph` instead of a
68    :class:`MultiGraph`.
69    """
70    import pydot
71
72    data = path.read()
73
74    # List of one or more "pydot.Dot" instances deserialized from this file.
75    P_list = pydot.graph_from_dot_data(data)
76
77    # Convert only the first such instance into a NetworkX graph.
78    return from_pydot(P_list[0])
79
80
81def from_pydot(P):
82    """Returns a NetworkX graph from a Pydot graph.
83
84    Parameters
85    ----------
86    P : Pydot graph
87      A graph created with Pydot
88
89    Returns
90    -------
91    G : NetworkX multigraph
92        A MultiGraph or MultiDiGraph.
93
94    Examples
95    --------
96    >>> K5 = nx.complete_graph(5)
97    >>> A = nx.nx_pydot.to_pydot(K5)
98    >>> G = nx.nx_pydot.from_pydot(A)  # return MultiGraph
99
100    # make a Graph instead of MultiGraph
101    >>> G = nx.Graph(nx.nx_pydot.from_pydot(A))
102
103    """
104    if P.get_strict(None):  # pydot bug: get_strict() shouldn't take argument
105        multiedges = False
106    else:
107        multiedges = True
108
109    if P.get_type() == "graph":  # undirected
110        if multiedges:
111            N = nx.MultiGraph()
112        else:
113            N = nx.Graph()
114    else:
115        if multiedges:
116            N = nx.MultiDiGraph()
117        else:
118            N = nx.DiGraph()
119
120    # assign defaults
121    name = P.get_name().strip('"')
122    if name != "":
123        N.name = name
124
125    # add nodes, attributes to N.node_attr
126    for p in P.get_node_list():
127        n = p.get_name().strip('"')
128        if n in ("node", "graph", "edge"):
129            continue
130        N.add_node(n, **p.get_attributes())
131
132    # add edges
133    for e in P.get_edge_list():
134        u = e.get_source()
135        v = e.get_destination()
136        attr = e.get_attributes()
137        s = []
138        d = []
139
140        if isinstance(u, str):
141            s.append(u.strip('"'))
142        else:
143            for unodes in u["nodes"]:
144                s.append(unodes.strip('"'))
145
146        if isinstance(v, str):
147            d.append(v.strip('"'))
148        else:
149            for vnodes in v["nodes"]:
150                d.append(vnodes.strip('"'))
151
152        for source_node in s:
153            for destination_node in d:
154                N.add_edge(source_node, destination_node, **attr)
155
156    # add default attributes for graph, nodes, edges
157    pattr = P.get_attributes()
158    if pattr:
159        N.graph["graph"] = pattr
160    try:
161        N.graph["node"] = P.get_node_defaults()[0]
162    except (IndexError, TypeError):
163        pass  # N.graph['node']={}
164    try:
165        N.graph["edge"] = P.get_edge_defaults()[0]
166    except (IndexError, TypeError):
167        pass  # N.graph['edge']={}
168    return N
169
170
171def to_pydot(N):
172    """Returns a pydot graph from a NetworkX graph N.
173
174    Parameters
175    ----------
176    N : NetworkX graph
177      A graph created with NetworkX
178
179    Examples
180    --------
181    >>> K5 = nx.complete_graph(5)
182    >>> P = nx.nx_pydot.to_pydot(K5)
183
184    Notes
185    -----
186
187    """
188    import pydot
189
190    # set Graphviz graph type
191    if N.is_directed():
192        graph_type = "digraph"
193    else:
194        graph_type = "graph"
195    strict = nx.number_of_selfloops(N) == 0 and not N.is_multigraph()
196
197    name = N.name
198    graph_defaults = N.graph.get("graph", {})
199    if name == "":
200        P = pydot.Dot("", graph_type=graph_type, strict=strict, **graph_defaults)
201    else:
202        P = pydot.Dot(
203            f'"{name}"', graph_type=graph_type, strict=strict, **graph_defaults
204        )
205    try:
206        P.set_node_defaults(**N.graph["node"])
207    except KeyError:
208        pass
209    try:
210        P.set_edge_defaults(**N.graph["edge"])
211    except KeyError:
212        pass
213
214    for n, nodedata in N.nodes(data=True):
215        str_nodedata = {k: str(v) for k, v in nodedata.items()}
216        p = pydot.Node(str(n), **str_nodedata)
217        P.add_node(p)
218
219    if N.is_multigraph():
220        for u, v, key, edgedata in N.edges(data=True, keys=True):
221            str_edgedata = {k: str(v) for k, v in edgedata.items() if k != "key"}
222            edge = pydot.Edge(str(u), str(v), key=str(key), **str_edgedata)
223            P.add_edge(edge)
224
225    else:
226        for u, v, edgedata in N.edges(data=True):
227            str_edgedata = {k: str(v) for k, v in edgedata.items()}
228            edge = pydot.Edge(str(u), str(v), **str_edgedata)
229            P.add_edge(edge)
230    return P
231
232
233def graphviz_layout(G, prog="neato", root=None):
234    """Create node positions using Pydot and Graphviz.
235
236    Returns a dictionary of positions keyed by node.
237
238    Parameters
239    ----------
240    G : NetworkX Graph
241        The graph for which the layout is computed.
242    prog : string (default: 'neato')
243        The name of the GraphViz program to use for layout.
244        Options depend on GraphViz version but may include:
245        'dot', 'twopi', 'fdp', 'sfdp', 'circo'
246    root : Node from G or None (default: None)
247        The node of G from which to start some layout algorithms.
248
249    Returns
250    -------
251      Dictionary of (x, y) positions keyed by node.
252
253    Examples
254    --------
255    >>> G = nx.complete_graph(4)
256    >>> pos = nx.nx_pydot.graphviz_layout(G)
257    >>> pos = nx.nx_pydot.graphviz_layout(G, prog="dot")
258
259    Notes
260    -----
261    This is a wrapper for pydot_layout.
262    """
263    return pydot_layout(G=G, prog=prog, root=root)
264
265
266def pydot_layout(G, prog="neato", root=None):
267    """Create node positions using :mod:`pydot` and Graphviz.
268
269    Parameters
270    ----------
271    G : Graph
272        NetworkX graph to be laid out.
273    prog : string  (default: 'neato')
274        Name of the GraphViz command to use for layout.
275        Options depend on GraphViz version but may include:
276        'dot', 'twopi', 'fdp', 'sfdp', 'circo'
277    root : Node from G or None (default: None)
278        The node of G from which to start some layout algorithms.
279
280    Returns
281    -------
282    dict
283        Dictionary of positions keyed by node.
284
285    Examples
286    --------
287    >>> G = nx.complete_graph(4)
288    >>> pos = nx.nx_pydot.pydot_layout(G)
289    >>> pos = nx.nx_pydot.pydot_layout(G, prog="dot")
290
291    Notes
292    -----
293    If you use complex node objects, they may have the same string
294    representation and GraphViz could treat them as the same node.
295    The layout may assign both nodes a single location. See Issue #1568
296    If this occurs in your case, consider relabeling the nodes just
297    for the layout computation using something similar to::
298
299        H = nx.convert_node_labels_to_integers(G, label_attribute='node_label')
300        H_layout = nx.nx_pydot.pydot_layout(G, prog='dot')
301        G_layout = {H.nodes[n]['node_label']: p for n, p in H_layout.items()}
302
303    """
304    import pydot
305
306    P = to_pydot(G)
307    if root is not None:
308        P.set("root", str(root))
309
310    # List of low-level bytes comprising a string in the dot language converted
311    # from the passed graph with the passed external GraphViz command.
312    D_bytes = P.create_dot(prog=prog)
313
314    # Unique string decoded from these bytes with the preferred locale encoding
315    D = str(D_bytes, encoding=getpreferredencoding())
316
317    if D == "":  # no data returned
318        print(f"Graphviz layout with {prog} failed")
319        print()
320        print("To debug what happened try:")
321        print("P = nx.nx_pydot.to_pydot(G)")
322        print('P.write_dot("file.dot")')
323        print(f"And then run {prog} on file.dot")
324        return
325
326    # List of one or more "pydot.Dot" instances deserialized from this string.
327    Q_list = pydot.graph_from_dot_data(D)
328    assert len(Q_list) == 1
329
330    # The first and only such instance, as guaranteed by the above assertion.
331    Q = Q_list[0]
332
333    node_pos = {}
334    for n in G.nodes():
335        pydot_node = pydot.Node(str(n)).get_name()
336        node = Q.get_node(pydot_node)
337
338        if isinstance(node, list):
339            node = node[0]
340        pos = node.get_pos()[1:-1]  # strip leading and trailing double quotes
341        if pos is not None:
342            xx, yy = pos.split(",")
343            node_pos[n] = (float(xx), float(yy))
344    return node_pos
345