mirror of
https://github.com/flatcar/scripts.git
synced 2025-08-08 13:36:58 +02:00
d0b0503276
This will be re-used by a separate tool Review URL: http://codereview.chromium.org/2859039
117 lines
3.4 KiB
Python
117 lines
3.4 KiB
Python
#!/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()
|