testing strings along with binary strings

8 files changed, 192 insertions, 10 deletions

diff --git a/Manifest.txt b/Manifest.txt
index 9225889..e1b1ee3 100644
--- a/Manifest.txt
+++ b/Manifest.txt
@@ -32,10 +32,12 @@ lib/psych/visitors/visitor.rb
 lib/psych/visitors/yast_builder.rb
 test/psych/test_parser.rb
 test/psych/test_serialize_subclasses.rb
+test/psych/test_to_yaml_properties.rb
 test/psych/test_tree_builder.rb
 test/test_psych.rb
 test/test_scalar_scanner.rb
 test/visitors/test_emitter.rb
 test/visitors/test_to_ruby.rb
 test/visitors/test_yast_builder.rb
+test/yaml/test_string.rb
 test/yaml/test_yaml.rb
diff --git a/lib/psych.rb b/lib/psych.rb
index 9250a2c..df28ba0 100644
--- a/lib/psych.rb
+++ b/lib/psych.rb
@@ -42,8 +42,10 @@ module Psych
 
   ###
   # Dump object +o+ to a YAML string
-  def self.dump o
-    o.to_yaml
+  def self.dump o, options = {}
+    visitor = Psych::Visitors::YASTBuilder.new options
+    visitor.accept o
+    visitor.tree.to_yaml
   end
 
   ###
diff --git a/lib/psych/ruby.rb b/lib/psych/ruby.rb
index f564565..47e4421 100644
--- a/lib/psych/ruby.rb
+++ b/lib/psych/ruby.rb
@@ -13,8 +13,6 @@ class Object
   include Psych::Visitable
 
   def to_yaml options = {}
-    visitor = Psych::Visitors::YASTBuilder.new options
-    visitor.accept self
-    visitor.tree.to_yaml
+    Psych.dump self, options
   end
 end
diff --git a/lib/psych/visitors/to_ruby.rb b/lib/psych/visitors/to_ruby.rb
index dd8805b..f91537f 100644
--- a/lib/psych/visitors/to_ruby.rb
+++ b/lib/psych/visitors/to_ruby.rb
@@ -75,6 +75,14 @@ module Psych
 
       def visit_Psych_Nodes_Mapping o
         case o.tag
+        when '!str', 'tag:yaml.org,2002:str'
+          members = Hash[*o.children.map { |c| accept c }]
+          string = members.delete 'str'
+
+          members.each do |k,v|
+            string.instance_variable_set k, v
+          end
+          string
         when /!ruby\/struct:?(.*)?$/
           klassname = $1
           members = o.children.map { |c| accept c }
diff --git a/lib/psych/visitors/yast_builder.rb b/lib/psych/visitors/yast_builder.rb
index 6c5f25f..6e2295c 100644
--- a/lib/psych/visitors/yast_builder.rb
+++ b/lib/psych/visitors/yast_builder.rb
@@ -26,7 +26,13 @@ module Psych
         klass = o.class == Object ? nil : o.class.name
         tag = ['!ruby/object', klass].compact.join(':')
         @stack.push append Nodes::Mapping.new(nil, tag, false)
-        o.instance_variables.each do |iv|
+        if o.respond_to? :to_yaml_properties
+          ivars = o.to_yaml_properties
+        else
+          ivars = o.instance_variables
+        end
+
+        ivars.each do |iv|
           accept iv.to_s.sub(/^@/, '')
           accept o.instance_variable_get(iv)
         end
@@ -41,7 +47,14 @@ module Psych
           accept member
           accept o[member]
         end
-        o.instance_variables.each do |iv|
+
+        if o.respond_to? :to_yaml_properties
+          ivars = o.to_yaml_properties
+        else
+          ivars = o.instance_variables
+        end
+
+        ivars.each do |iv|
           accept iv.to_s.sub(/^@/, '')
           accept o.instance_variable_get(iv)
         end
@@ -114,10 +127,43 @@ module Psych
       end
 
       def visit_String o
-        quote = ScalarScanner.new(o).tokenize.first != :SCALAR
+        plain = false
+        quote = false
+
+        if o.index("\x00") || o.count("^ -~\t\r\n").fdiv(o.length) > 0.3
+          str   = [o].pack('m').chomp
+          tag   = '!binary'
+        else
+          str   = o
+          tag   = nil
+          quote = ScalarScanner.new(o).tokenize.first != :SCALAR
+          plain = !quote
+        end
 
-        scalar = Nodes::Scalar.new(o, nil, nil, !quote, quote)
-        @stack.last.children << scalar
+
+        if o.respond_to? :to_yaml_properties
+          ivars = o.to_yaml_properties
+        else
+          ivars = o.instance_variables
+        end
+
+        scalar = Nodes::Scalar.new str, nil, tag, plain, quote
+
+        if ivars.empty?
+          append scalar
+        else
+          mapping = append Nodes::Mapping.new(nil, '!str', false)
+
+          mapping.children << Nodes::Scalar.new('str')
+          mapping.children << scalar
+
+          @stack.push mapping
+          ivars.each do |iv|
+            mapping.children << Nodes::Scalar.new(":#{iv}")
+            accept o.instance_variable_get(iv)
+          end
+          @stack.pop
+        end
       end
 
       def visit_Class o
diff --git a/test/psych/test_to_yaml_properties.rb b/test/psych/test_to_yaml_properties.rb
new file mode 100644
index 0000000..ff35fcd
--- /dev/null
+++ b/test/psych/test_to_yaml_properties.rb
@@ -0,0 +1,64 @@
+require 'minitest/autorun'
+require 'psych'
+
+module Psych
+  class TestToYamlProperties < MiniTest::Unit::TestCase
+    class Foo
+      attr_accessor :a, :b, :c
+      def initialize
+        @a = 1
+        @b = 2
+        @c = 3
+      end
+
+      def to_yaml_properties
+        [:@a, :@b]
+      end
+    end
+
+    def test_object_dump_yaml_properties
+      foo = Psych.load(Psych.dump(Foo.new))
+      assert_equal 1, foo.a
+      assert_equal 2, foo.b
+      assert_nil foo.c
+    end
+
+    class Bar < Struct.new(:foo, :bar)
+      attr_reader :baz
+      def initialize *args
+        super
+        @baz = 'hello'
+      end
+
+      def to_yaml_properties
+        []
+      end
+    end
+
+    def test_struct_dump_yaml_properties
+      bar = Psych.load(Psych.dump(Bar.new('a', 'b')))
+      assert_equal 'a', bar.foo
+      assert_equal 'b', bar.bar
+      assert_nil bar.baz
+    end
+
+    def test_string_dump
+      string = "okonomiyaki"
+      class << string
+        def to_yaml_properties
+          [:@tastes]
+        end
+      end
+
+      string.instance_variable_set(:@tastes, 'delicious')
+      v = Psych.load Psych.dump string
+      assert_equal 'delicious', v.instance_variable_get(:@tastes)
+    end
+
+    def test_string_load_syck
+      str = Psych.load("--- !str \nstr: okonomiyaki\n:@tastes: delicious\n")
+      assert_equal 'okonomiyaki', str
+      assert_equal 'delicious', str.instance_variable_get(:@tastes)
+    end
+  end
+end
diff --git a/test/visitors/test_emitter.rb b/test/visitors/test_emitter.rb
index e79a7b5..f440678 100644
--- a/test/visitors/test_emitter.rb
+++ b/test/visitors/test_emitter.rb
@@ -59,6 +59,21 @@ module Psych
         assert_equal @io.string, s.to_yaml
       end
 
+      def test_scalar_with_tag
+        s       = Nodes::Stream.new
+        doc     = Nodes::Document.new
+        scalar  = Nodes::Scalar.new 'hello world', nil, '!str', false, false, 5
+
+        doc.children << scalar
+        s.children << doc
+
+        @visitor.accept s
+
+        assert_match(/str/, @io.string)
+        assert_match(/hello/, @io.string)
+        assert_equal @io.string, s.to_yaml
+      end
+
       def test_sequence
         s       = Nodes::Stream.new
         doc     = Nodes::Document.new
diff --git a/test/yaml/test_string.rb b/test/yaml/test_string.rb
new file mode 100644
index 0000000..7857672
--- /dev/null
+++ b/test/yaml/test_string.rb
@@ -0,0 +1,47 @@
+require 'test/unit'
+require 'psych'
+
+YAML = Psych
+
+module YAML
+  class TestString < Test::Unit::TestCase
+    def test_binary_string_null
+      string = "\x00"
+      yml = YAML.dump string
+      assert_match(/binary/, yml)
+      assert_equal string, YAML.load(yml)
+    end
+
+    def test_binary_string
+      string = binary_string
+      yml = YAML.dump string
+      assert_match(/binary/, yml)
+      assert_equal string, YAML.load(yml)
+    end
+
+    def test_non_binary_string
+      string = binary_string(0.29)
+      yml = YAML.dump string
+      refute_match(/binary/, yml)
+      assert_equal string, YAML.load(yml)
+    end
+
+    def test_string_with_ivars
+      food = "is delicious"
+      ivar = "on rock and roll"
+      food.instance_variable_set(:@we_built_this_city, ivar)
+
+      str = YAML.load YAML.dump food
+      assert_equal ivar, food.instance_variable_get(:@we_built_this_city)
+    end
+
+    def binary_string percentage = 0.31, length = 100
+      string = ''
+      (percentage * length).to_i.times do |i|
+        string << "\b"
+      end
+      string << 'a' * (length - string.length)
+      string
+    end
+  end
+end