src/third_party/chromite/lib/dot_helper.py

   1 # Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """Helper functions for building graphs with dot."""
   6
   7 from __future__ import print_function
   8
   9 from chromite.lib import cros_build_lib
  10 from chromite.lib import osutils
  11
  12
  13 class Subgraph(object):
  14   """A subgraph in dot. Contains nodes, arcs, and other subgraphs."""
  15
  16   _valid_ranks = set(['source', 'sink', 'same', 'min', 'max', None])
  17
  18   def __init__(self, rank=None):
  19     self._rank = rank
  20     self._nodes = []
  21     self._subgraphs = []
  22     self._arcs = set()
  23     self._rank = None
  24
  25   def AddNode(self, node_id, name=None, color=None, href=None):
  26     """Adds a node to the subgraph."""
  27     tags = {}
  28     if name:
  29       tags['label'] = name
  30     if color:
  31       tags['color'] = color
  32       tags['fontcolor'] = color
  33     if href:
  34       tags['href'] = href
  35     self._nodes.append({'id': node_id, 'tags': tags})
  36
  37   def AddSubgraph(self, subgraph):
  38     """Adds a subgraph to the subgraph."""
  39     self._subgraphs.append(subgraph)
  40
  41   def AddNewSubgraph(self, rank=None):
  42     """Adds a new subgraph to the subgraph. The new subgraph is returned."""
  43     subgraph = Subgraph(rank)
  44     self.AddSubgraph(subgraph)
  45     return subgraph
  46
  47   def AddArc(self, node_from, node_to):
  48     """Adds an arc between two nodes."""
  49     self._arcs.add((node_from, node_to))
  50
  51   def _GenNodes(self):
  52     """Generates the code for all the nodes."""
  53     lines = []
  54     for node in self._nodes:
  55       tags = ['%s="%s"' % (k, v) for (k, v) in node['tags'].iteritems()]
  56       lines.append('"%s" [%s];' % (node['id'], ', '.join(tags)))
  57     return lines
  58
  59   def _GenSubgraphs(self):
  60     """Generates the code for all the subgraphs contained in this subgraph."""
  61     lines = []
  62     for subgraph in self._subgraphs:
  63       lines += subgraph.Gen()
  64     return lines
  65
  66   def _GenArcs(self):
  67     """Generates the code for all the arcs."""
  68     lines = []
  69     for node_from, node_to in self._arcs:
  70       lines.append('"%s" -> "%s";' % (node_from, node_to))
  71     return lines
  72
  73   def _GenInner(self):
  74     """Generates the code for the inner contents of the subgraph."""
  75     lines = []
  76     if self._rank:
  77       lines.append('rank=%s;' % self._rank)
  78     lines += self._GenSubgraphs()
  79     lines += self._GenNodes()
  80     lines += self._GenArcs()
  81     return lines
  82
  83   def Gen(self):
  84     """Generates the code for the subgraph."""
  85     return ['subgraph {'] + self._GenInner() + ['}']
  86
  87
  88 class Graph(Subgraph):
  89   """A top-level graph in dot. It's basically a subgraph with a name."""
  90
  91   def __init__(self, name):
  92     Subgraph.__init__(self)
  93     self._name = name
  94
  95   def Gen(self):
  96     """Generates the code for the graph."""
  97     return ['digraph "%s" {' % self._name,
  98             'graph [name="%s"];' % self._name] + self._GenInner() + ['}']
  99
 100
 101 def GenerateImage(lines, filename, out_format='svg', save_dot_filename=None):
 102   """Generates the image by calling dot on the input lines."""
 103   data = '\n'.join(lines)
 104   cros_build_lib.RunCommand(['dot', '-T%s' % out_format, '-o', filename],
 105                             input=data)
 106
 107   if save_dot_filename:
 108     osutils.WriteFile(save_dot_filename, data)