lib/rgl/dot.rb in rgl-0.5.3 vs lib/rgl/dot.rb in rgl-0.5.4

@@ -31,25 +31,33 @@
     def to_dot_graph(params = {})
       params['name'] ||= self.class.name.gsub(/:/, '_')
       fontsize       = params['fontsize'] ? params['fontsize'] : '8'
       graph          = (directed? ? DOT::Digraph : DOT::Graph).new(params)
       edge_class     = directed? ? DOT::DirectedEdge : DOT::Edge
+      vertex_options = params['vertex'] || {}
+      edge_options   = params['edge'] || {}
 
       each_vertex do |v|
-        graph << DOT::Node.new(
-            'name'     => vertex_id(v),
-            'fontsize' => fontsize,
-            'label'    => vertex_label(v)
-        )
+        default_vertex_options =  {
+          'name'     => vertex_id(v),
+          'fontsize' => fontsize,
+          'label'    => vertex_label(v)
+        }
+        each_vertex_options = default_vertex_options.merge(vertex_options)
+        vertex_options.each{|option, val| each_vertex_options[option] = val.call(v) if val.is_a?(Proc)}
+        graph << DOT::Node.new(each_vertex_options)
       end
 
       each_edge do |u, v|
-        graph << edge_class.new(
-            'from'     => vertex_id(u),
-            'to'       => vertex_id(v),
-            'fontsize' => fontsize
-        )
+        default_edge_options = {
+          'from'     => vertex_id(u),
+          'to'       => vertex_id(v),
+          'fontsize' => fontsize
+        }
+        each_edge_options = default_edge_options.merge(edge_options)
+        edge_options.each{|option, val| each_edge_options[option] = val.call(u, v) if val.is_a?(Proc)}
+        graph << edge_class.new(each_edge_options)
       end
 
       graph
     end
 
@@ -73,15 +81,15 @@
     end
 
     # Use dot[http://www.graphviz.org] to create a graphical representation of
     # the graph. Returns the filename of the graphics file.
     #
-    def write_to_graphic_file(fmt='png', dotfile="graph")
+    def write_to_graphic_file(fmt='png', dotfile="graph", options={})
       src = dotfile + ".dot"
       dot = dotfile + "." + fmt
 
       File.open(src, 'w') do |f|
-        f << self.to_dot_graph.to_s << "\n"
+        f << self.to_dot_graph(options).to_s << "\n"
       end
 
       unless system("dot -T#{fmt} #{src} -o #{dot}")
         raise "Error executing dot. Did you install GraphViz?"
       end