graph 1.2.0 → 2.0.0

data/lib/graph.rb

@@ -1,130 +1,165 @@
 #!/usr/local/bin/ruby -w
 
 ##
-# Graph is a type of hash that outputs in graphviz's dot format.
+# Graph models directed graphs and subgraphs and outputs in graphviz's
+# dot format.
+
+class Graph
+  VERSION = "2.0.0" # :nodoc:
+
+  LIGHT_COLORS = %w(gray lightblue lightcyan lightgray lightpink
+                    lightslategray lightsteelblue white)
+
+  # WTF -- can't be %w() because of a bug in rcov
+  BOLD_COLORS = ["black", "brown", "mediumblue", "blueviolet",
+                 "orange", "magenta", "darkgreen", "maroon",
+                 "violetred", "purple", "greenyellow", "deeppink",
+                 "midnightblue", "firebrick", "darkturquoise",
+                 "mediumspringgreen", "chartreuse", "navy",
+                 "lightseagreen", "chocolate", "lawngreen", "green",
+                 "indigo", "darkgoldenrod", "darkviolet", "red",
+                 "springgreen", "saddlebrown", "mediumvioletred",
+                 "goldenrod", "tomato", "cyan", "forestgreen",
+                 "darkorchid", "crimson", "coral", "deepskyblue",
+                 "seagreen", "peru", "turquoise", "orangered",
+                 "dodgerblue", "sienna", "limegreen", "royalblue",
+                 "darkorange", "blue"]
+
+  SHAPES = %w(Mcircle Mdiamond Msquare box box3d circle component
+              diamond doublecircle doubleoctagon egg ellipse folder
+              hexagon house invhouse invtrapezium invtriangle none
+              note octagon parallelogram pentagon plaintext point
+              polygon rect rectangle septagon square tab trapezium
+              triangle tripleoctagon)
+
+  STYLES = %w(dashed dotted solid invis bold filled diagonals rounded)
+
+  STYLES.each do |name|
+    define_method(name) { style name }
+  end
+
+  (BOLD_COLORS + LIGHT_COLORS).each do |name|
+    define_method(name) { color name }
+  end
 
-class Graph < Hash
-  VERSION = '1.2.0' # :nodoc:
+  SHAPES.each do |name|
+    define_method(name.downcase) { shape name }
+  end
 
   ##
-  # A Hash of arrays of attributes for each node. Eg:
-  #
-  #   graph.attribs["a"] << "color = red"
-  #
-  # Will color node "a" red.
+  # A parent graph, if any. Only used for subgraphs.
 
-  attr_reader :attribs
+  attr_accessor :graph
 
   ##
-  # An array of the order of traversal / definition of the nodes.
-  #
-  # You (generally) should leave this alone.
+  # The name of the graph. Optional for graphs and subgraphs. Prefix
+  # the name of a subgraph with "cluster" for subgraph that is boxed.
 
-  attr_reader :order
+  attr_accessor :name
 
   ##
-  # An array of attributes to add to the front of the graph source. Eg:
-  #
-  #   graph.prefix << "ratio = 1.5"
+  # Global attributes for edges in this graph.
 
-  attr_reader :prefix
+  attr_reader :edge_attribs
 
   ##
-  # A Hash of a Hashes of Arrays of attributes for an edge. Eg:
-  #
-  #   graph.edge["a"]["b"] << "color = blue"
-  #
-  # Will color the edge between a and b blue.
+  # The hash of hashes of edges in this graph. Use #[] or #node to create edges.
 
-  attr_reader :edge
+  attr_reader :edges
 
-  def []= key, val # :nodoc:
-    @order << key unless self.has_key? key
-    super
-  end
+  ##
+  # Global attributes for this graph.
 
-  def clear # :nodoc:
-    super
-    @prefix.clear
-    @order.clear
-    @attribs.clear
-    @edge.clear
-  end
+  attr_reader :graph_attribs
 
   ##
-  # Return all nodes (aka the set of all keys and values).
+  # Global attributes for nodes in this graph.
 
-  def nodes
-    (keys + values).flatten.uniq
-  end
+  attr_reader :node_attribs
 
   ##
-  # A convenience method to set every node to be a box. Especially
-  # good for longer text nodes.
+  # The hash of nodes in this graph. Use #[] or #node to create nodes.
 
-  def boxes
-    prefix << "node [ shape = box ]"
+  attr_reader :nodes
+
+  ##
+  # An array of subgraphs.
+
+  attr_reader :subgraphs
+
+  ##
+  # Creates a new graph object. Optional name and parent graph are
+  # available. Also takes an optional block for DSL-like use.
+
+  def initialize name = nil, graph = nil, &block
+    @name = name
+    @graph = graph
+    graph << self if graph
+    @nodes  = Hash.new { |h,k| h[k] = Node.new self, k }
+    @edges  = Hash.new { |h,k|
+      h[k] = Hash.new { |h2, k2| h2[k2] = Edge.new self, self[k], self[k2] }
+    }
+    @graph_attribs = []
+    @node_attribs  = []
+    @edge_attribs  = []
+    @subgraphs     = []
+
+    instance_eval(&block) if block
   end
 
   ##
-  # A convenience method to set an attribute on all nodes.
+  # Push a subgraph into the current graph. Sets the subgraph's graph to self.
 
-  def global_attrib attrib
-    nodes.each do |key|
-      attribs[key] << attrib
-    end
+  def << subgraph
+    subgraphs << subgraph
+    subgraph.graph = self
   end
 
   ##
-  # Returns a Hash with a count of the outgoing edges for each node.
+  # Access a node by name
 
-  def counts
-    result = Hash.new 0
-    each_pair do |from, to|
-      result[from] += 1
-    end
-    result
+  def [] name
+    nodes[name]
   end
 
-  def delete key # :nodoc:
-    @order.delete key
-    # TODO: prolly needs to go clean up attribs
-    super
+  ##
+  # A convenience method to set the global node attributes to use boxes.
+
+  def boxes
+    node_attribs << shape("box")
   end
 
   ##
-  # Overrides Hash#each_pair to go over each _edge_. Eg:
-  #
-  #   g["a"] << "b"
-  #   g["a"] << "b"
-  #   g.each_pair { |from, to| ... }
-  #
-  # goes over a -> b *twice*.
+  # Shortcut method to create a new color Attribute instance.
 
-  def each_pair
-    @order.each do |from|
-      self[from].each do |to|
-        yield from, to
-      end
-    end
+  def color color
+    Attribute.new "color = #{color}"
   end
 
   ##
-  # Deletes nodes that have less than minimum number of outgoing edges.
+  # Shortcut method to create a new colorscheme Attribute instance.
 
-  def filter_size minimum
-    counts.each do |node, count|
-      next unless count < minimum
-      delete node
-    end
+  def colorscheme name
+    Attribute.new "colorscheme = #{name}"
   end
 
-  def initialize # :nodoc:
-    super { |h,k| h[k] = [] }
-    @prefix  = []
-    @order   = []
-    @attribs = Hash.new { |h,k| h[k] = [] }
-    @edge    = Hash.new { |h,k| h[k] = Hash.new { |h2,k2| h2[k2] = [] } }
+  ##
+  # Define one or more edges.
+  #
+  #   edge "a", "b", "c", ...
+  #
+  # is equivalent to:
+  #
+  #   edge "a", "b"
+  #   edge "b", "c"
+  #   ...
+
+  def edge(*names)
+    last = nil
+    names.each_cons(2) do |from, to|
+      last = self[from][to]
+    end
+    last
   end
 
   ##
@@ -132,76 +167,261 @@ class Graph < Hash
 
   def invert
     result = self.class.new
-    each_pair do |from, to|
-      result[to] << from
+    edges.each do |from, h|
+      h.each do |to, edge|
+        result[to][from]
+      end
     end
     result
   end
 
   ##
-  # Returns a list of keys sorted by number of outgoing edges.
+  # Shortcut method to create a new fillcolor Attribute instance.
 
-  def keys_by_count
-    counts.sort_by { |key, count| -count }.map {|key, count| key }
+  def fillcolor n
+    Attribute.new "fillcolor = #{n}"
   end
 
   ##
-  # Specify the orientation of the graph. Defaults to the graphviz default "TB".
+  # Shortcut method to create a new font Attribute instance. You can
+  # pass in both the name and an optional font size.
 
-  def orient dir = "TB"
-    prefix << "rankdir = #{dir}"
+  def font name, size=nil
+    Attribute.new "fontname = #{name.inspect}" +
+      (size ? ", fontsize = #{size}" : "")
   end
 
   ##
-  # Really just an alias for #orient but with "LR" as the default value.
+  # Shortcut method to set the graph's label. Usually used with subgraphs.
 
-  def rotate dir = "LR"
-    orient dir
+  def label name
+    graph_attribs << "label = \"#{name}\""
   end
 
   ##
-  # Remove all duplicate edges.
+  # Access a node by name, supplying an optional label
 
-  def normalize
-    each do |k,v|
-      v.uniq!
-    end
+  def node name, label = nil
+    n = nodes[name]
+    n.label label if label
+    n
+  end
+
+  ##
+  # Shortcut method to specify the orientation of the graph. Defaults
+  # to the graphviz default "TB".
+
+  def orient dir = "TB"
+    graph_attribs << "rankdir = #{dir}"
+  end
+
+  ##
+  # Shortcut method to specify the orientation of the graph. Defaults to "LR".
+
+  def rotate dir = "LR"
+    orient dir
   end
 
   ##
   # Saves out both a dot file to path and an image for the specified type.
   # Specify type as nil to skip exporting an image.
 
-  def save path, type="png"
+  def save path, type = nil
     File.open "#{path}.dot", "w" do |f|
       f.write self.to_s
     end
     system "dot -T#{type} #{path}.dot > #{path}.#{type}" if type
   end
 
+  ##
+  # Shortcut method to create a new shape Attribute instance.
+
+  def shape shape
+    Attribute.new "shape = #{shape}"
+  end
+
+  ##
+  # Shortcut method to create a new style Attribute instance.
+
+  def style name
+    Attribute.new "style = #{name}"
+  end
+
+  ##
+  # Shortcut method to create a subgraph in the current graph. Use
+  # with the top-level +digraph+ method in block form for a graph DSL.
+
+  def subgraph name = nil, &block
+    Graph.new name, self, &block
+  end
+
   ##
   # Outputs a graphviz graph.
 
   def to_s
     result = []
-    result << "digraph absent"
+
+    type = graph ? "subgraph" : "digraph"
+    result << "#{type} #{name}"
     result << "  {"
 
-    @prefix.each do |line|
+    graph_attribs.each do |line|
       result << "    #{line};"
     end
 
-    @attribs.sort.each do |node, attribs|
-      result << "    #{node.inspect} [ #{attribs.join ','} ];"
+    unless node_attribs.empty? then
+      result << "    node [ #{node_attribs.join(", ")} ];"
+    end
+
+    unless edge_attribs.empty? then
+      result << "    edge [ #{edge_attribs.join(", ")} ];"
     end
 
-    each_pair do |from, to|
-      edge = @edge[from][to].join ", "
-      edge = " [ #{edge} ]" unless edge.empty?
-      result << "    #{from.inspect} -> #{to.inspect}#{edge};"
+    subgraphs.each do |line|
+      result << "    #{line};"
+    end
+
+    nodes.each do |name, node|
+      result << "    #{node};" if graph or not node.attributes.empty?
+    end
+
+    edges.each do |from, deps|
+      deps.each do |to, edge|
+        result << "    #{edge};"
+      end
     end
 
     result << "  }"
     result.join "\n"
   end
+
+  ##
+  # An attribute for a graph, node, or edge. Really just a composable
+  # string (via #+) with a convenience method #<< that allows you to
+  # "paint" nodes and edges with this attribute.
+
+  class Attribute < Struct.new :attr
+    ##
+    # "Paint" graphs, nodes, and edges with this attribute.
+    #
+    #   red << node1 << node2 << node3
+    #
+    # is the same as:
+    #
+    #   node1.attributes << red
+    #   node2.attributes << red
+    #   node3.attributes << red
+
+    def << thing
+      thing.attributes << self
+      self
+    end
+
+    ##
+    # Returns the attribute in string form.
+
+    alias :to_s :attr
+
+    ##
+    # Compose a new attribute from two existing attributes:
+    #
+    #   bad_nodes = red + filled + diamond
+
+    def + style
+      Attribute.new "#{self}, #{style}"
+    end
+  end
+
+  ##
+  # An edge in a graph.
+
+  class Edge < Struct.new :graph, :from, :to, :attributes
+
+    ##
+    # Create a new edge in +graph+ from +from+ to +to+.
+
+    def initialize graph, from, to
+      super graph, from, to, []
+    end
+
+    ##
+    # Shortcut method to set the label attribute on an edge.
+
+    def label name
+      attributes << "label = \"#{name}\""
+      self
+    end
+
+    ##
+    # Returns the edge in dot syntax.
+
+    def to_s
+      fromto = "%p -> %p" % [from.name, to.name]
+      if attributes.empty? then
+        fromto
+      else
+        "%-20s [ %-20s ]" % [fromto, attributes.join(',')]
+      end
+    end
+  end
+
+  ##
+  # Nodes in the graph.
+
+  class Node < Struct.new :graph, :name, :attributes
+
+    ##
+    # Create a new Node. Takes a parent graph and a name.
+
+    def initialize graph, name
+      super graph, name, []
+    end
+
+    ##
+    # Shortcut method to set the node's label.
+
+    def label name
+      attributes << "label = #{name.inspect}"
+    end
+
+    ##
+    # Create a new node with +name+ and an edge between them pointing
+    # from self to the new node.
+
+    def >> name
+      self[name] # creates node and edge
+      self
+    end
+
+    alias :"<<" :">>"
+
+    ##
+    # Returns the edge between self and +dep_name+.
+
+    def [] dep_name
+      graph.edges[name][dep_name]
+    end
+
+    ##
+    # Returns the node in dot syntax.
+
+    def to_s
+      if attributes.empty? then
+        "#{name.inspect}"
+      else
+        "%-20p [ %-20s ]" % [name, attributes.join(',')]
+      end
+    end
+  end
+end
+
+##
+# Convenience method to create a new graph. Used for DSL-style:
+#
+#   g = digraph do
+#     edge "a", "b", "c"
+#   end
+
+def digraph name = nil, &block
+  Graph.new name, &block
 end