Index: dot_helper.py |
diff --git a/dot_helper.py b/dot_helper.py |
new file mode 100644 |
index 0000000000000000000000000000000000000000..d22f901f92fb19a2cb630793e8e52f89e14657a6 |
--- /dev/null |
+++ b/dot_helper.py |
@@ -0,0 +1,116 @@ |
+#!/usr/bin/python |
+# Copyright (c) 2010 The Chromium OS Authors. All rights reserved. |
+# Use of this source code is governed by a BSD-style license that can be |
+# found in the LICENSE file. |
+ |
+"""Helper functions for building graphs with dot.""" |
+ |
+import subprocess |
+ |
+class Subgraph(object): |
+ """A subgraph in dot. Contains nodes, arcs, and other subgraphs.""" |
+ |
+ _valid_ranks = set(['source', 'sink', 'same', 'min', 'max', None]) |
+ |
+ def __init__(self, rank=None): |
+ self.SetRank(rank) |
+ self._nodes = [] |
+ self._subgraphs = [] |
+ self._arcs = set() |
+ |
+ def SetRank(self, rank): |
+ """Sets the rank for the nodes in the graph. |
+ |
+ Can be one of 'source', 'sink', 'same', 'min', 'max' or None. See dot |
+ documentation at http://www.graphviz.org/Documentation.php for exact |
+ semantics.""" |
+ assert(rank in self._valid_ranks) |
+ self._rank = rank |
+ |
+ def AddNode(self, id, name=None, color=None, href=None): |
+ """Adds a node to the subgraph.""" |
+ tags = {} |
+ if name: |
+ tags['label'] = name |
+ if color: |
+ tags['color'] = color |
+ tags['fontcolor'] = color |
+ if href: |
+ tags['href'] = href |
+ self._nodes.append({'id': id, 'tags': tags}) |
+ |
+ def AddSubgraph(self, subgraph): |
+ """Adds a subgraph to the subgraph.""" |
+ self._subgraphs.append(subgraph) |
+ |
+ def AddNewSubgraph(self, rank=None): |
+ """Adds a new subgraph to the subgraph. The new subgraph is returned.""" |
+ subgraph = Subgraph(rank) |
+ self.AddSubgraph(subgraph) |
+ return subgraph |
+ |
+ def AddArc(self, node_from, node_to): |
+ """Adds an arc between two nodes.""" |
+ self._arcs.add((node_from, node_to)) |
+ |
+ def _GenNodes(self): |
+ """Generates the code for all the nodes.""" |
+ lines = [] |
+ for node in self._nodes: |
+ tags = ['%s="%s"' % (k, v) for (k, v) in node['tags'].iteritems()] |
+ lines.append('"%s" [%s];' % (node['id'], ', '.join(tags))) |
+ return lines |
+ |
+ def _GenSubgraphs(self): |
+ """Generates the code for all the subgraphs contained in this subgraph.""" |
+ lines = [] |
+ for subgraph in self._subgraphs: |
+ lines += subgraph.Gen(); |
+ return lines |
+ |
+ def _GenArcs(self): |
+ """Generates the code for all the arcs.""" |
+ lines = [] |
+ for node_from, node_to in self._arcs: |
+ lines.append('"%s" -> "%s";' % (node_from, node_to)) |
+ return lines |
+ |
+ def _GenInner(self): |
+ """Generates the code for the inner contents of the subgraph.""" |
+ lines = [] |
+ if self._rank: |
+ lines.append('rank=%s;' % self._rank) |
+ lines += self._GenSubgraphs() |
+ lines += self._GenNodes() |
+ lines += self._GenArcs() |
+ return lines |
+ |
+ def Gen(self): |
+ """Generates the code for the subgraph.""" |
+ return ['subgraph {'] + self._GenInner() + ['}'] |
+ |
+ |
+class Graph(Subgraph): |
+ """A top-level graph in dot. It's basically a subgraph with a name.""" |
+ |
+ def __init__(self, name): |
+ Subgraph.__init__(self) |
+ self._name = name |
+ |
+ def Gen(self): |
+ """Generates the code for the graph.""" |
+ return ['digraph "%s" {' % self._name, |
+ 'graph [name="%s"];' % self._name] + self._GenInner() + ['}'] |
+ |
+ |
+def GenerateImage(lines, filename, format='svg', save_dot_filename=None): |
+ """Generates the image by calling dot on the input lines.""" |
+ data = '\n'.join(lines) |
+ proc = subprocess.Popen(['dot', '-T' + format, '-o' + filename], |
+ stdin=subprocess.PIPE) |
+ proc.communicate(data) |
+ |
+ if save_dot_filename: |
+ file = open(save_dot_filename, 'w') |
+ file.write(data) |
+ file.close() |