classes can be emitted as scalar and map

5 files changed, 103 insertions, 39 deletions

diff --git a/lib/psych.rb b/lib/psych.rb
index cebc427..0c1ad22 100644
--- a/lib/psych.rb
+++ b/lib/psych.rb
@@ -7,6 +7,7 @@ require 'psych/parser'
 require 'psych/omap'
 require 'psych/set'
 require 'psych/coder'
+require 'psych/yaml'
 
 ###
 # = Overview
@@ -156,6 +157,19 @@ module Psych
   def self.add_domain_type domain, type_tag, &block
     @domain_types[type_tag] = [domain, block]
   end
-  class << self; attr_accessor :domain_types; end
+
+  @load_tags = {}
+  @dump_tags = {}
+  def self.add_tag tag, klass
+    @load_tags[tag] = klass
+    @dump_tags[klass] = tag
+  end
+
+  class << self
+    attr_accessor :load_tags
+    attr_accessor :dump_tags
+    attr_accessor :domain_types
+  end
+
   # :startdoc:
 end
diff --git a/lib/psych/coder.rb b/lib/psych/coder.rb
index 88c7322..05ee886 100644
--- a/lib/psych/coder.rb
+++ b/lib/psych/coder.rb
@@ -1,33 +1,50 @@
 module Psych
-  class Coder < ::Hash
-    def initialize map, h = nil
-      super()
-      merge!(h) if h
-      @map = map
-    end
+  ###
+  # If an object defines +encode_with+, then an instance of Psych::Coder will
+  # passed to the method when the object is being serialized.  The Coder
+  # automatically assumes a Psych::Nodes::Mapping is being emitted.  Other
+  # objects like Sequence and Scalar may be emitted if +seq=+ or +scalar=+ are
+  # called, respectively.
+  class Coder
+    attr_accessor :tag, :style, :implicit
+    attr_reader   :type, :map, :scalar, :seq
 
-    def tag= tag
-      @map.tag = tag
+    def initialize tag
+      @map        = {}
+      @seq        = []
+      @implicit   = false
+      @type       = :map
+      @tag        = tag
+      @style      = Psych::Nodes::Mapping::BLOCK
+      @scalar     = nil
     end
 
-    def tag
-      @map.tag
+    # Emit a scalar with +value+
+    def scalar= value
+      @type   = :scalar
+      @scalar = value
     end
 
-    def style= style
-      @map.style = style
+    # Emit a map with +value+
+    def map= map
+      @type = :map
+      @map  = map
     end
 
-    def style
-      @map.style
+    def []= k, v
+      @type = :map
+      @map[k] = v
     end
 
-    def implicit= implicity
-      @map.implicit = implicity
+    def [] k
+      @type = :map
+      @map[k]
     end
 
-    def implicit
-      @map.implicit
+    # Emit a sequence of +list+
+    def seq= list
+      @type = :seq
+      @seq  = list
     end
   end
 end
diff --git a/lib/psych/visitors/to_ruby.rb b/lib/psych/visitors/to_ruby.rb
index 2ec0449..e080ec6 100644
--- a/lib/psych/visitors/to_ruby.rb
+++ b/lib/psych/visitors/to_ruby.rb
@@ -85,6 +85,8 @@ module Psych
       end
 
       def visit_Psych_Nodes_Mapping o
+        return revive(Psych.load_tags[o.tag], o) if Psych.load_tags[o.tag]
+
         case o.tag
         when '!str', 'tag:yaml.org,2002:str'
           members = Hash[*o.children.map { |c| accept c }]
@@ -144,9 +146,7 @@ module Psych
 
         when /^!ruby\/object:?(.*)?$/
           name = $1 || 'Object'
-          h = Hash[*o.children.map { |c| accept c }]
-          s = (resolve_class(name) || Object).allocate
-          init_with(s, h, o)
+          revive((resolve_class(name) || Object), o)
         else
           hash = {}
           @st[o.anchor] = hash if o.anchor
@@ -170,9 +170,17 @@ module Psych
       end
 
       private
+      def revive klass, node
+        s = klass.allocate
+        h = Hash[*node.children.map { |c| accept c }]
+        init_with(s, h, node)
+      end
+
       def init_with o, h, node
         if o.respond_to?(:init_with)
-          o.init_with Psych::Coder.new(node, h)
+          c = Psych::Coder.new(node.tag)
+          c.map = h
+          o.init_with c
         else
           h.each { |k,v| o.instance_variable_set(:"@#{k}", v) }
         end
diff --git a/lib/psych/visitors/yaml_tree.rb b/lib/psych/visitors/yaml_tree.rb
index d27b1ec..120a78d 100644
--- a/lib/psych/visitors/yaml_tree.rb
+++ b/lib/psych/visitors/yaml_tree.rb
@@ -29,7 +29,11 @@ module Psych
           return append Nodes::Alias.new target.object_id.to_s
         end
 
-        send(@dispatch_cache[target.class], target)
+        if target.respond_to?(:encode_with)
+          dump_coder target
+        else
+          send(@dispatch_cache[target.class], target)
+        end
       end
 
       def visit_Psych_Omap o
@@ -42,8 +46,11 @@ module Psych
       end
 
       def visit_Object o
-        klass = o.class == Object ? nil : o.class.name
-        tag = ['!ruby/object', klass].compact.join(':')
+        tag = Psych.dump_tags[o.class]
+        unless tag
+          klass = o.class == Object ? nil : o.class.name
+          tag   = ['!ruby/object', klass].compact.join(':')
+        end
 
         map = append Nodes::Mapping.new(nil, tag, false)
 
@@ -233,23 +240,37 @@ module Psych
         yaml_obj
       end
 
-      def dump_ivars target, map
-        if target.respond_to?(:encode_with)
-          coder = Psych::Coder.new(map)
-          target.encode_with(coder)
-          coder.each do |k,v|
+      def dump_coder o
+        tag = Psych.dump_tags[o.class]
+        unless tag
+          klass = o.class == Object ? nil : o.class.name
+          tag   = ['!ruby/object', klass].compact.join(':')
+        end
+
+        c = Psych::Coder.new(tag)
+        o.encode_with(c)
+        case c.type
+        when :scalar
+          append Nodes::Scalar.new(c.scalar, nil, c.tag)
+        when :map
+          map = append Nodes::Mapping.new(nil, c.tag, c.implicit, c.style)
+          @stack.push map
+          c.map.each do |k,v|
             map.children << Nodes::Scalar.new(k)
             accept v
           end
-        else
-          ivars = target.respond_to?(:to_yaml_properties) ?
-            target.to_yaml_properties :
-            target.instance_variables
+          @stack.pop
+        end
+      end
 
-          ivars.each do |iv|
-            map.children << Nodes::Scalar.new("#{iv.to_s.sub(/^@/, '')}")
-            accept target.instance_variable_get(iv)
-          end
+      def dump_ivars target, map
+        ivars = target.respond_to?(:to_yaml_properties) ?
+          target.to_yaml_properties :
+          target.instance_variables
+
+        ivars.each do |iv|
+          map.children << Nodes::Scalar.new("#{iv.to_s.sub(/^@/, '')}")
+          accept target.instance_variable_get(iv)
         end
       end
     end
diff --git a/lib/psych/yaml.rb b/lib/psych/yaml.rb
index 9ac007d..db40a03 100644
--- a/lib/psych/yaml.rb
+++ b/lib/psych/yaml.rb
@@ -12,6 +12,10 @@ YAML = Psych
 end
 
 class Object
+  def self.yaml_tag name
+    Psych.add_tag(name, self)
+  end
+
   def to_yaml options = {}
     Psych.dump self, options
   end