lib/end_state/graph.rb in end_state-0.2.0 vs lib/end_state/graph.rb in end_state-0.3.0

@@ -1,20 +1,27 @@
 module EndState
   class Graph < GraphViz
-    attr_reader :machine, :nodes
+    attr_reader :machine, :nodes, :event_labels
 
-    def initialize(machine)
+    def initialize(machine, event_labels=true)
       @machine = machine
       @nodes = {}
+      @event_labels = event_labels
       super machine.name.to_sym
     end
 
     def draw
       machine.transitions.keys.each do |t|
         left, right = t.to_a.flatten
         nodes[left] ||= add_node(left.to_s)
         nodes[right] ||= add_node(right.to_s)
-        add_edge nodes[left], nodes[right]
+        edge = add_edge nodes[left], nodes[right]
+        if event_labels
+          event = machine.events.detect do |event, transition|
+            transition.include? t
+          end
+          edge[:label] = event.first.to_s if event
+        end
       end
       self
     end
   end
 end
\ No newline at end of file