Split graph() function; Fix #96

circuits · Jan 24, 2017 · ec9f5b9 · ec9f5b9
1 parent 25442fc
commit ec9f5b9
Showing 1 changed file with 36 additions and 28 deletions.
diff --git a/circuits/tools/__init__.py b/circuits/tools/__init__.py
@@ -60,7 +60,7 @@ def kill(x):
         x.unregister()
 
 
-def graph(x, name=None):
+def graph_dot(x, name=None):
     """Display a directed graph of the Component structure of x
 
     :param x: A Component or Manager to graph
@@ -77,49 +77,57 @@ def graph(x, name=None):
     pygraphviz = tryimport("pygraphviz")
     plt = tryimport("matplotlib.pyplot", "pyplot")
 
-    if networkx is not None and pygraphviz is not None and plt is not None:
-        graph_edges = []
-        for (u, v, d) in edges(x):
-            graph_edges.append((u.name, v.name, float(d)))
+    if not all([networkx, pygraphviz, plt]):
+        return
 
-        g = networkx.DiGraph()
-        g.add_weighted_edges_from(graph_edges)
+    graph_edges = []
+    for (u, v, d) in edges(x):
+        graph_edges.append((u.name, v.name, float(d)))
 
-        elarge = [(u, v) for (u, v, d) in g.edges(data=True)
-                  if d["weight"] > 3.0]
-        esmall = [(u, v) for (u, v, d) in g.edges(data=True)
-                  if d["weight"] <= 3.0]
+    g = networkx.DiGraph()
+    g.add_weighted_edges_from(graph_edges)
 
-        pos = networkx.spring_layout(g)  # positions for all nodes
+    elarge = [(u, v) for (u, v, d) in g.edges(data=True) if d["weight"] > 3.0]
+    esmall = [(u, v) for (u, v, d) in g.edges(data=True) if d["weight"] <= 3.0]
 
-        # nodes
-        networkx.draw_networkx_nodes(g, pos, node_size=700)
+    pos = networkx.spring_layout(g)  # positions for all nodes
 
-        # edges
-        networkx.draw_networkx_edges(g, pos, edgelist=elarge, width=1)
-        networkx.draw_networkx_edges(
-            g, pos, edgelist=esmall, width=1,
-            alpha=0.5, edge_color="b", style="dashed"
-        )
+    # nodes
+    networkx.draw_networkx_nodes(g, pos, node_size=700)
 
-        # labels
-        networkx.draw_networkx_labels(
-            g, pos, font_size=10, font_family="sans-serif"
-        )
+    # edges
+    networkx.draw_networkx_edges(g, pos, edgelist=elarge, width=1)
+    networkx.draw_networkx_edges(
+        g, pos, edgelist=esmall, width=1,
+        alpha=0.5, edge_color="b", style="dashed"
+    )
+
+    # labels
+    networkx.draw_networkx_labels(
+        g, pos, font_size=10, font_family="sans-serif"
+    )
 
-        plt.axis("off")
+    plt.axis("off")
 
-        plt.savefig("{0:s}.png".format(name or x.name))
-        networkx.write_dot(g, "{0:s}.dot".format(name or x.name))
+    plt.savefig("{0:s}.png".format(name or x.name))
+    networkx.write_dot(g, "{0:s}.dot".format(name or x.name))
 
-        plt.clf()
+    plt.clf()
+    return g
 
+
+def graph_ascii(x):
     def printer(d, x):
         return "%s* %s" % (" " * d, x)
 
     return "\n".join(walk(x, printer))
 
 
+def graph(x, name=None):
+    graph_dot(x, name)
+    return graph_ascii(x)
+
+
 def inspect(x):
     """Display an inspection report of the Component or Manager x