Create simple GraphViz dot code from a graph builder API in Python.

dot_graph_builder.py

"""Python GraphViz dot generator. Create simple GraphViz dot code from a graph
builder API.

The goal is to hook this up with a GUI at some point.

Future Improvements
  * Potentially subprocess the graphviz generator to allow this to spit out
    graphviz SVG or PNG files!
  * UI that calls this API when you click certain elements.
  
I don't like the verbosity of dot syntax and I don't need all its complex features.
Most of my block diagrams involve directed edges with blocks of varying shape and
color.
"""
from __future__ import print_function

from collections import defaultdict


class GVNode(object):
  """Node in the graphviz graph."""

  def __init__(self,
               name,
               color='black',
               shape='ellipse',
               text=None,
               subtext=None):
    self.name = name
    self.color = color
    self.shape = shape
    self.text = text
    self.subtext = subtext

  def to_str(self):
    """Get the GraphViz dot textual representation of this node."""
    if not self.text:
      self.text = self.name
    return ('"{0}" [color={1}, shape={2}, label=<{3}<br/><font '
            'point-size="10">{4}</font>>];').format(
                self.name, self.color, self.shape, self.text, self.subtext)


class GVDirectedEdge(object):
  """Edge in the graphviz graph."""

  def __init__(self, src, dest, color):
    self.src = src
    self.dest = dest
    self.color = color

  def to_str(self):
    return '"{}"->"{}" [color="{}"];'.format(self.src, self.dest, self.color)


class GVGraph(object):
  """Representation of the overall graph."""

  def __init__(self, name='new_graph'):
    self.name = name
    self.edge_matrix = defaultdict(lambda: defaultdict(str))
    self.nodes = {}

  def add_node(self,
               name,
               color='black',
               shape='ellipse',
               text=None,
               subtext=None):

    self.nodes[name] = GVNode(name, color, shape, text, subtext)

  def add_edge(self, src_name, dest_name, color='black'):
    self.edge_matrix[src_name][dest_name] = GVDirectedEdge(
        src_name, dest_name, color)

  def to_str(self):
    raw_dot = 'digraph {} {{'.format(self.name)

    # Add each node in the graph.
    for node_name in self.nodes:
      raw_dot += self.nodes[node_name].to_str()

    # Add each edge in the graph.
    for edge_src in self.edge_matrix:
      for edge_dest in self.edge_matrix[edge_src]:
        raw_dot += self.edge_matrix[edge_src][edge_dest].to_str()

    raw_dot += '}'
    return raw_dot

  def get_node(self, node_name):
    return self.nodes[node_name]


if __name__ == '__main__':
  # Here's some example usage!
  # Create a graph.
  graph = GVGraph()
  graph.add_node('Node1')
  n = graph.get_node('Node1')
  n.fill = 'gold'
  n.outline = 'blue'
  n.text = 'This is some text.'
  n.subtext = 'This is some subtext.'
  graph.add_node('Node2')
  graph.add_edge('Node1', 'Node2')
  print(graph.to_str())