simplify pydot imports, use testing.utils routines

imports now just import pydotplus
use assert_graphs_equal
fix from_pydot to no longer create '' as attributes
    'name', 'graph', 'node' and 'edge'

2 files changed, 40 insertions, 54 deletions

diff --git a/networkx/drawing/nx_pydot.py b/networkx/drawing/nx_pydot.py
index 4a936966..b3fc414a 100644
--- a/networkx/drawing/nx_pydot.py
+++ b/networkx/drawing/nx_pydot.py
@@ -3,7 +3,7 @@
 Pydot
 *****
 
-Import and export NetworkX graphs in Graphviz dot format using pydot.
+Import and export NetworkX graphs in Graphviz dot format using pydotplus.
 
 Either this module or nx_pygraphviz can be used to interface with graphviz.
 
@@ -59,12 +59,9 @@ def read_dot(path):
     -----
     Use G=nx.Graph(nx.read_dot(path)) to return a Graph instead of a MultiGraph.
     """
-    try:
-        import pydotplus as pydot
-    except ImportError:
-        raise ImportError("pydotplus required for read_dot()")
+    import pydotplus
     data = path.read()
-    P = pydot.graph_from_dot_data(data)
+    P = pydotplus.graph_from_dot_data(data)
     return from_pydot(P)
 
 def from_pydot(P):
@@ -95,18 +92,19 @@ def from_pydot(P):
 
     if P.get_type()=='graph': # undirected
         if multiedges:
-            create_using=nx.MultiGraph()
+            N = nx.MultiGraph()
         else:
-            create_using=nx.Graph()
+            N = nx.Graph()
     else:
         if multiedges:
-            create_using=nx.MultiDiGraph()
+            N = nx.MultiDiGraph()
         else:
-            create_using=nx.DiGraph()
+            N = nx.DiGraph()
 
     # assign defaults
-    N=nx.empty_graph(0,create_using)
-    N.name=P.get_name()
+    name=P.get_name().strip('"')
+    if name != '':
+        N.name = name
 
     # add nodes, attributes to N.node_attr
     for p in P.get_node_list():
@@ -140,15 +138,17 @@ def from_pydot(P):
                 N.add_edge(source_node,destination_node,**attr)
 
     # add default attributes for graph, nodes, edges
-    N.graph['graph']=P.get_attributes()
+    pattr = P.get_attributes()
+    if pattr:
+        N.graph['graph'] = pattr
     try:
         N.graph['node']=P.get_node_defaults()[0]
     except:# IndexError,TypeError:
-        N.graph['node']={}
+        pass #N.graph['node']={}
     try:
         N.graph['edge']=P.get_edge_defaults()[0]
     except:# IndexError,TypeError:
-        N.graph['edge']={}
+        pass #N.graph['edge']={}
     return N
 
 def to_pydot(N, strict=True):
@@ -168,10 +168,7 @@ def to_pydot(N, strict=True):
     -----
 
     """
-    try:
-        import pydotplus as pydot
-    except ImportError:
-        raise ImportError("pydotplus required for to_pydot()")
+    import pydotplus
     # set Graphviz graph type
     if N.is_directed():
         graph_type='digraph'
@@ -179,12 +176,13 @@ def to_pydot(N, strict=True):
         graph_type='graph'
     strict=N.number_of_selfloops()==0 and not N.is_multigraph()
 
-    name = N.graph.get('name')
+    name = N.name
     graph_defaults=N.graph.get('graph',{})
-    if name is None:
-        P = pydot.Dot(graph_type=graph_type,strict=strict,**graph_defaults)
+    if name is '':
+        P = pydotplus.Dot('', graph_type=graph_type, strict=strict,
+                      **graph_defaults)
     else:
-        P = pydot.Dot('"%s"'%name,graph_type=graph_type,strict=strict,
+        P = pydotplus.Dot('"%s"'%name, graph_type=graph_type, strict=strict,
                       **graph_defaults)
     try:
         P.set_node_defaults(**N.graph['node'])
@@ -197,19 +195,20 @@ def to_pydot(N, strict=True):
 
     for n,nodedata in N.nodes(data=True):
         str_nodedata=dict((k,make_str(v)) for k,v in nodedata.items())
-        p=pydot.Node(make_str(n),**str_nodedata)
+        p=pydotplus.Node(make_str(n),**str_nodedata)
         P.add_node(p)
 
     if N.is_multigraph():
         for u,v,key,edgedata in N.edges(data=True,keys=True):
             str_edgedata=dict((k,make_str(v)) for k,v in edgedata.items())
-            edge=pydot.Edge(make_str(u),make_str(v),key=make_str(key),**str_edgedata)
+            edge=pydotplus.Edge(make_str(u), make_str(v),
+                    key=make_str(key), **str_edgedata)
             P.add_edge(edge)
 
     else:
         for u,v,edgedata in N.edges(data=True):
             str_edgedata=dict((k,make_str(v)) for k,v in edgedata.items())
-            edge=pydot.Edge(make_str(u),make_str(v),**str_edgedata)
+            edge=pydotplus.Edge(make_str(u),make_str(v),**str_edgedata)
             P.add_edge(edge)
     return P
 
@@ -256,10 +255,7 @@ def pydot_layout(G,prog='neato',root=None, **kwds):
     >>> pos=nx.pydot_layout(G)
     >>> pos=nx.pydot_layout(G,prog='dot')
     """
-    try:
-        import pydotplus as pydot
-    except ImportError:
-        raise ImportError("pydotplus required for to_pydot()")
+    import pydotplus
     P=to_pydot(G)
     if root is not None :
         P.set("root",make_str(root))
@@ -275,11 +271,11 @@ def pydot_layout(G,prog='neato',root=None, **kwds):
         print("And then run %s on file.dot"%(prog))
         return
 
-    Q=pydot.graph_from_dot_data(D)
+    Q=pydotplus.graph_from_dot_data(D)
 
     node_pos={}
     for n in G.nodes():
-        pydot_node = pydot.Node(make_str(n)).get_name()
+        pydot_node = pydotplus.Node(make_str(n)).get_name()
         node=Q.get_node(pydot_node)
 
         if isinstance(node,list):
@@ -294,6 +290,6 @@ def pydot_layout(G,prog='neato',root=None, **kwds):
 def setup_module(module):
     from nose import SkipTest
     try:
-        import pydotplus as pydot
+        import pydotplus
     except ImportError:
         raise SkipTest("pydotplus not available")
diff --git a/networkx/drawing/tests/test_pydot.py b/networkx/drawing/tests/test_pydot.py
index 72e67cbd..2ae26c38 100644
--- a/networkx/drawing/tests/test_pydot.py
+++ b/networkx/drawing/tests/test_pydot.py
@@ -8,42 +8,31 @@ from nose import SkipTest
 from nose.tools import assert_true
 
 import networkx as nx
+from networkx.testing import assert_graphs_equal
 
 class TestPydot(object):
     @classmethod
     def setupClass(cls):
-        global pydot
+        global pydotplus
         try:
-            import pydotplus as pydot
+            import pydotplus
         except ImportError:
             raise SkipTest('pydotplus not available.')
 
-    def build_graph(self, G):
+    def pydot_checks(self, G):
         G.add_edge('A','B')
         G.add_edge('A','C')
         G.add_edge('B','C')
         G.add_edge('A','D')
         G.add_node('E')
-        return G, nx.to_pydot(G)
-
-    def assert_equal(self, G1, G2):
-        assert_true( sorted(G1.nodes())==sorted(G2.nodes()) )
-        if G1.is_directed():
-            assert_true(sorted(G1.edges()) == sorted(G2.edges()))
-        else:
-            e1 = sorted(map(sorted, G1.edges()))
-            e2 = sorted(map(sorted, G2.edges()))
-            assert_true(e1 == e2)
-
-    def pydot_checks(self, G):
-        H, P = self.build_graph(G)
-        G2 = H.__class__(nx.from_pydot(P))
-        self.assert_equal(H, G2)
+        P = nx.to_pydot(G)
+        G2 = G.__class__(nx.from_pydot(P))
+        assert_graphs_equal(G, G2)
 
         fname = tempfile.mktemp()
         assert_true( P.write_raw(fname) )
 
-        Pin = pydot.graph_from_dot_file(fname)
+        Pin = pydotplus.graph_from_dot_file(fname)
 
         n1 = sorted([p.get_name() for p in P.get_node_list()])
         n2 = sorted([p.get_name() for p in Pin.get_node_list()])
@@ -54,8 +43,9 @@ class TestPydot(object):
         assert_true( sorted(e1)==sorted(e2) )
 
         Hin = nx.drawing.nx_pydot.read_dot(fname)
-        Hin = H.__class__(Hin)
-        self.assert_equal(H, Hin)
+        Hin = G.__class__(Hin)
+        assert_graphs_equal(G, Hin)
+
 #        os.unlink(fname)