1 files changed, 90 insertions, 0 deletions
diff --git a/steuermann/dot.py b/steuermann/dot.py
new file mode 100644
index 0000000..e4d7a1f
--- /dev/null
+++ b/steuermann/dot.py
@@ -0,0 +1,90 @@
+import pygraphviz as gv
+import sys, os
+import nodes
+
+
+'''
+This script produces a DOT file of the nodes in one or more *.sm files
+
+Requires graphviz and pygraphviz
+'''
+
+
+if len(sys.argv) < 2:
+    print 'ERROR - missing argument(s)'
+    sys.exit(1)
+
+sm_files = sys.argv[1:]
+
+# get a command-tree of the *.sm files
+tree = nodes.read_file_list(sm_files)
+
+# do some things to get a list of only direct predecessors
+N = []
+for node in tree.node_index.values():
+    node.parents = [p.name for p in node.predecessors]
+    N.append(node)
+
+hosts = []
+for node in N:
+    if node.host not in hosts:
+        hosts.append(node.host)
+
+extras = {}
+for node in N:
+    if node.name not in extras.keys():
+        extras[node.name] = []
+    for s in node.predecessors:
+        for p in s.parents:
+            if p in node.parents:
+                extras[node.name].append(p)
+
+for k, v in extras.items():
+    extras[k] = list(set(v))
+
+for i, node in enumerate(N):
+    node.direct_parents = []
+    for s in node.predecessors:
+        if s.name not in extras[node.name]:
+            node.direct_parents.append(s)
+    N[i] = node
+
+
+# these are just some random colors (enough for 13 hosts)
+colors = [
+    'red',
+    'black',
+    'purple',
+    'lawngreen',
+    'tomato',
+    'teal',
+    'brown',
+    'orange',
+    'blue',
+    'midnightblue',
+    'lightseagreen',
+    'goldenrod',
+    'black'
+]
+
+# assign host colors
+host_colors = {}
+for host in hosts:
+    host_colors[host] = colors.pop(0)
+
+
+# make the graph
+graph = gv.AGraph(directed=True)
+
+for node in N:
+    graph.add_node(node.name, color = host_colors[node.host])
+
+for node in N:
+    for s in node.direct_parents:
+        graph.add_edge(s.name, node.name)
+
+
+# output to $HOME/sm_steps.dot
+home = os.environ['HOME']
+graph.write(os.path.join(home, 'sm_steps.dot'))
+