lib/psych/visitors/yaml_tree.rb in psych-1.3.4 vs lib/psych/visitors/yaml_tree.rb in psych-2.0.0

@@ -1,25 +1,60 @@
+require 'psych/tree_builder'
+require 'psych/scalar_scanner'
+require 'psych/class_loader'
+
 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
+      class Registrar # :nodoc:
+        def initialize
+          @obj_to_id   = {}
+          @obj_to_node = {}
+          @counter     = 0
+        end
+
+        def register target, node
+          @obj_to_node[target.object_id] = node
+        end
+
+        def key? target
+          @obj_to_node.key? target.object_id
+        end
+
+        def id_for target
+          @obj_to_id[target.object_id] ||= (@counter += 1)
+        end
+
+        def node_for target
+          @obj_to_node[target.object_id]
+        end
+      end
+
       attr_reader :started, :finished
       alias :finished? :finished
       alias :started? :started
 
-      def initialize options = {}, emitter = TreeBuilder.new, ss = ScalarScanner.new
+      def self.create options = {}, emitter = nil
+        emitter      ||= TreeBuilder.new
+        class_loader = ClassLoader.new
+        ss           = ScalarScanner.new class_loader
+        new(emitter, ss, options)
+      end
+
+      def initialize emitter, ss, options
         super()
         @started  = false
         @finished = false
         @emitter  = emitter
-        @st       = {}
+        @st       = Registrar.new
         @ss       = ss
         @options  = options
         @coders   = []
 
         @dispatch_cache = Hash.new do |h,klass|
@@ -45,10 +80,11 @@
         end
       end
 
       def tree
         finish unless finished?
+        @emitter.root
       end
 
       def push object
         start unless started?
         version = []
@@ -63,19 +99,19 @@
           version = [1,1]
         end if @options.key? :version
 
         @emitter.start_document version, [], false
         accept object
-        @emitter.end_document
+        @emitter.end_document !@emitter.streaming?
       end
       alias :<< :push
 
       def accept target
         # return any aliases we find
-        if @st.key? target.object_id
-          oid         = target.object_id
-          node        = @st[oid]
+        if @st.key? target
+          oid         = @st.id_for target
+          node        = @st.node_for target
           anchor      = oid.to_s
           node.anchor = anchor
           return @emitter.alias anchor
         end
 
@@ -218,31 +254,37 @@
       def visit_BigDecimal o
         @emitter.scalar o._dump, nil, '!ruby/object:BigDecimal', false, false, Nodes::Scalar::ANY
       end
 
       def binary? string
-        string.encoding == Encoding::ASCII_8BIT ||
+        (string.encoding == Encoding::ASCII_8BIT && !string.ascii_only?) ||
           string.index("\x00") ||
-          string.count("\x00-\x7F", "^ -~\t\r\n").fdiv(string.length) > 0.3
+          string.count("\x00-\x7F", "^ -~\t\r\n").fdiv(string.length) > 0.3 ||
+          string.class != String
       end
       private :binary?
 
       def visit_String o
-        plain = false
-        quote = false
-        style = Nodes::Scalar::ANY
+        plain = true
+        quote = true
+        style = Nodes::Scalar::PLAIN
+        tag   = nil
+        str   = o
 
         if binary?(o)
           str   = [o].pack('m').chomp
           tag   = '!binary' # FIXME: change to below when syck is removed
           #tag   = 'tag:yaml.org,2002:binary'
           style = Nodes::Scalar::LITERAL
+          plain = false
+          quote = false
+        elsif o =~ /\n/
+          style = Nodes::Scalar::LITERAL
         else
-          str   = o
-          tag   = nil
-          quote = !(String === @ss.tokenize(o))
-          plain = !quote
+          unless String === @ss.tokenize(o)
+            style = Nodes::Scalar::SINGLE_QUOTED
+          end
         end
 
         ivars = find_ivars o
 
         if ivars.empty?
@@ -400,11 +442,11 @@
 
         target.instance_variables
       end
 
       def register target, yaml_obj
-        @st[target.object_id] = yaml_obj
+        @st.register target, yaml_obj
         yaml_obj
       end
 
       def dump_coder o
         @coders << o
@@ -430,10 +472,10 @@
           end
           @emitter.end_sequence
         when :map
           @emitter.start_mapping nil, c.tag, c.implicit, c.style
           c.map.each do |k,v|
-            @emitter.scalar k, nil, nil, true, false, Nodes::Scalar::ANY
+            accept k
             accept v
           end
           @emitter.end_mapping
         when :object
           accept c.object