Refactor the dot graph generation out of cros_generate_deps_graphs

This will be re-used by a separate tool

Review URL: http://codereview.chromium.org/2859039

cros_generate_deps_graphs

@@ -5,9 +5,9 @@
 
 """Generates pretty dependency graphs for Chrome OS packages."""
 
+import dot_helper
 import optparse
 import os
-import subprocess
 import sys
 
 
@@ -41,26 +41,16 @@ def GetOutputBaseName(node, options):
                           options.format)
 
 
-def GetNodeLines(node, options, color):
+def AddNodeToSubgraph(subgraph, node, options, color):
   """Gets the dot definition for a node."""
   name = node['full_name']
-  tags = ['label="%s (%s)"' % (name, node['action']),
-          'color="%s"' % color,
-          'fontcolor="%s"' % color]
+  href = None
   if options.link:
     filename = GetOutputBaseName(node, options)
-    tags.append('href="%s%s"' % (options.base_url, filename))
-  return ['"%s" [%s];' % (name, ', '.join(tags))]
+    href = '%s%s' % (options.base_url, filename)
+  subgraph.AddNode(name, name, color, href)
 
 
-def GetReverseDependencyArcLines(node, options):
-  """Gets the dot definitions for the arcs leading to a node."""
-  lines = []
-  name = node['full_name']
-  for j in node['rev_deps']:
-    lines.append('"%s" -> "%s";' % (j, name))
-  return lines
-
 
 def GenerateDotGraph(package, deps_map, options):
   """Generates the dot source for the dependency graph leading to a node.
@@ -73,38 +63,29 @@ def GenerateDotGraph(package, deps_map, options):
   # definitions
   emitted = set()
 
-  lines = ['digraph dep {',
-           'graph [name="%s"];' % package]
+  graph = dot_helper.Graph(package)
 
   # Add all the children if we want them, all of them in their own subgraph,
   # as a sink. Keep the arcs outside of the subgraph though (it generates
   # better layout).
-  has_children = False
+  children_subgraph = None
   if options.children and node['deps']:
-    has_children = True
-    lines += ['subgraph {',
-              'rank=sink;']
-    arc_lines = []
+    children_subgraph = graph.AddNewSubgraph('sink')
     for child in node['deps']:
       child_node = deps_map[child]
-      lines += GetNodeLines(child_node, options, CHILD_COLOR)
+      AddNodeToSubgraph(children_subgraph, child_node, options, CHILD_COLOR)
       emitted.add(child)
-      # If child is in the rev_deps, we'll get the arc later.
-      if not child in node['rev_deps']:
-        arc_lines.append('"%s" -> "%s";' % (package, child))
-    lines += ['}']
-    lines += arc_lines
+      graph.AddArc(package, child)
 
   # Add the package in its own subgraph. If we didn't have children, make it
   # a sink
-  lines += ['subgraph {']
-  if has_children:
-    lines += ['rank=same;']
+  if children_subgraph:
+    rank = 'same'
   else:
-    lines += ['rank=sink;']
-  lines += GetNodeLines(node, options, TARGET_COLOR)
+    rank = 'sink'
+  package_subgraph = graph.AddNewSubgraph(rank)
+  AddNodeToSubgraph(package_subgraph, node, options, TARGET_COLOR)
   emitted.add(package)
-  lines += ['}']
 
   # Add all the other nodes, as well as all the arcs.
   for dep in deps:
@@ -113,11 +94,11 @@ def GenerateDotGraph(package, deps_map, options):
       color = NORMAL_COLOR
       if dep_node['action'] == 'seed':
         color = SEED_COLOR
-      lines += GetNodeLines(dep_node, options, color)
-    lines += GetReverseDependencyArcLines(dep_node, options)
+      AddNodeToSubgraph(graph, dep_node, options, color)
+    for j in dep_node['rev_deps']:
+      graph.AddArc(j, dep)
 
-  lines += ['}']
-  return lines
+  return graph.Gen()
 
 
 def GenerateImages(input, options):
@@ -126,20 +107,15 @@ def GenerateImages(input, options):
 
   for package in deps_map:
     lines = GenerateDotGraph(package, deps_map, options)
-    data = '\n'.join(lines)
 
     filename = os.path.join(options.output_dir,
                             GetOutputBaseName(deps_map[package], options))
 
-    # Send the source to dot.
-    proc = subprocess.Popen(['dot', '-T' + options.format, '-o' + filename],
-                            stdin=subprocess.PIPE)
-    proc.communicate(data)
-
+    save_dot_filename = None
     if options.save_dot:
-      file = open(filename + '.dot', 'w')
-      file.write(data)
-      file.close()
+      save_dot_filename = filename + '.dot'
+
+    dot_helper.GenerateImage(lines, filename, options.format, save_dot_filename)
 
 
 def main():

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()