using factory methods for tree building

2 files changed, 130 insertions, 21 deletions

diff --git a/lib/psych/visitors/yaml_stream.rb b/lib/psych/visitors/yaml_stream.rb
new file mode 100644
index 0000000..48cd745
--- /dev/null
+++ b/lib/psych/visitors/yaml_stream.rb
@@ -0,0 +1,99 @@
+module Psych
+  module Visitors
+    class YAMLStream < Psych::Visitors::YAMLTree
+      class StreamEmitter
+        def initialize io
+          @handler = Psych::Emitter.new io
+          @stack   = []
+        end
+
+        def push node
+          @stack << node
+
+          case node
+          when Psych::Nodes::Mapping  then start_mapping node
+          when Psych::Nodes::Sequence then start_sequence node
+          when Psych::Nodes::Document then start_document node
+          when Psych::Nodes::Stream   then start_stream node
+          else
+            raise "Can't handle #{node}"
+          end
+        end
+        alias :<< :push
+
+        def pop
+          node = @stack.pop
+
+          case node
+          when Psych::Nodes::Mapping  then end_mapping node
+          when Psych::Nodes::Sequence then end_sequence node
+          when Psych::Nodes::Document then end_document node
+          when Psych::Nodes::Stream   then end_stream node
+          else
+            raise "Can't handle #{node}"
+          end
+        end
+
+        def start_stream o
+          @handler.start_stream o.encoding
+        end
+
+        def end_stream o
+          @handler.end_stream
+        end
+
+        def start_document o
+          @handler.start_document o.version, o.tag_directives, o.implicit
+        end
+
+        def end_document o
+          @handler.end_document o.implicit_end
+        end
+
+        def start_mapping o
+          @handler.start_mapping o.anchor, o.tag, o.implicit, o.style
+        end
+
+        def end_mapping o
+          @handler.end_mapping
+        end
+
+        def start_sequence o
+          @handler.start_sequence o.anchor, o.tag, o.implicit, o.style
+        end
+
+        def end_sequence o
+          @handler.end_sequence
+        end
+
+        def scalar value, anchor, tag, plain, quoted, style
+          @handler.scalar value, anchor, tag, plain, quoted, style
+        end
+      end
+
+      def initialize io
+        super()
+        @stack = StreamEmitter.new io
+        @stack << Psych::Nodes::Stream.new
+      end
+
+      def << object
+        @stack.push create_document
+        accept object
+        @stack.pop
+      end
+
+      def finish
+        @stack.pop
+      end
+
+      private
+      def append o; o end
+      def register o, yml; yml end
+
+      def create_scalar value, anchor = nil, tag = nil, plain = true, quoted = false, style = Nodes::Scalar::ANY
+        @stack.scalar value, anchor, tag, plain, quoted, style
+      end
+    end
+  end
+end
diff --git a/lib/psych/visitors/yaml_tree.rb b/lib/psych/visitors/yaml_tree.rb
index 608da6d..6f87b27 100644
--- a/lib/psych/visitors/yaml_tree.rb
+++ b/lib/psych/visitors/yaml_tree.rb
@@ -1,5 +1,12 @@
 module Psych
   module Visitors
+    ###
+    # YAMLTree builds a YAML ast given a ruby object.  For example:
+    #
+    #   builder = Psych::Visitors::YAMLTree.new
+    #   builder << { :foo => 'bar' }
+    #   builder.tree # => #<Psych::Nodes::Stream .. }
+    #
     class YAMLTree < Psych::Visitors::Visitor
       attr_reader :tree
 
@@ -75,23 +82,22 @@ module Psych
         register(o, map)
 
         @stack.push map
-        dump_ivars(o, map)
+        dump_ivars o
         @stack.pop
       end
 
       def visit_Struct o
         tag = ['!ruby/struct', o.class.name].compact.join(':')
 
-        map = register(o, create_mapping(nil, tag, false))
+        @stack.push append register(o, create_mapping(nil, tag, false))
 
-        @stack.push append map
 
         o.members.each do |member|
-          map.children <<  create_scalar("#{member}")
+          append create_scalar("#{member}")
           accept o[member]
         end
 
-        dump_ivars(o, map)
+        dump_ivars o
 
         @stack.pop
       end
@@ -99,20 +105,19 @@ module Psych
       def visit_Exception o
         tag = ['!ruby/exception', o.class.name].join ':'
 
-        map = append create_mapping(nil, tag, false)
+        @stack.push append create_mapping(nil, tag, false)
 
-        @stack.push map
 
         {
           'message'   => private_iv_get(o, 'mesg'),
           'backtrace' => private_iv_get(o, 'backtrace'),
         }.each do |k,v|
           next unless v
-          map.children << create_scalar(k)
+          append create_scalar(k)
           accept v
         end
 
-        dump_ivars(o, map)
+        dump_ivars o
 
         @stack.pop
       end
@@ -133,21 +138,26 @@ module Psych
       end
 
       def visit_Rational o
-        map = append create_mapping(nil, '!ruby/object:Rational', false)
+        @stack.push append create_mapping(nil, '!ruby/object:Rational', false)
+
         [
           'denominator', o.denominator.to_s,
           'numerator', o.numerator.to_s
         ].each do |m|
-          map.children << create_scalar(m)
+          append create_scalar m
         end
+
+        @stack.pop
       end
 
       def visit_Complex o
-        map = append create_mapping(nil, '!ruby/object:Complex', false)
+        @stack.push append create_mapping(nil, '!ruby/object:Complex', false)
 
         ['real', o.real.to_s, 'image', o.imag.to_s].each do |m|
-          map.children << create_scalar(m)
+          append create_scalar m
         end
+
+        @stack.pop
       end
 
       def visit_Integer o
@@ -191,13 +201,13 @@ module Psych
         if ivars.empty?
           append scalar
         else
-          mapping = append create_mapping(nil, '!str', false)
+          @stack.push append create_mapping(nil, '!str', false)
+
+          append create_scalar('str')
+          append scalar
 
-          mapping.children << create_scalar('str')
-          mapping.children << scalar
+          dump_ivars o
 
-          @stack.push mapping
-          dump_ivars o, mapping
           @stack.pop
         end
       end
@@ -300,18 +310,18 @@ module Psych
           map = append create_mapping(nil, c.tag, c.implicit, c.style)
           @stack.push map
           c.map.each do |k,v|
-            map.children << create_scalar(k)
+            append create_scalar k
             accept v
           end
           @stack.pop
         end
       end
 
-      def dump_ivars target, map
+      def dump_ivars target
         ivars = find_ivars target
 
         ivars.each do |iv|
-          map.children << create_scalar("#{iv.to_s.sub(/^@/, '')}")
+          append create_scalar("#{iv.to_s.sub(/^@/, '')}")
           accept target.instance_variable_get(iv)
         end
       end