From 2a4d9568f7d5d19c00231cf48eb855cc45ec3394 Mon Sep 17 00:00:00 2001
From: Sean Griffin <sean@thoughtbot.com>
Date: Wed, 17 Sep 2014 13:47:48 -0600
Subject: Allow dumping any BasicObject that defines #marshal_dump or
 #marshal_load

Additionally, it is assumed that these objects will have implemented
`respond_to?` and `class`, (e.g. SimpleDelegator). These assumptions are
reasonably inline with the assumptions of `Marshal` in regards to
`BasicObject`s.
---
 lib/psych/visitors/to_ruby.rb   | 15 ++++++++++++
 lib/psych/visitors/yaml_tree.rb | 14 +++++++++++
 test/psych/test_marshalable.rb  | 54 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 83 insertions(+)
 create mode 100644 test/psych/test_marshalable.rb

diff --git a/lib/psych/visitors/to_ruby.rb b/lib/psych/visitors/to_ruby.rb
index e74d5d4..e696ebd 100644
--- a/lib/psych/visitors/to_ruby.rb
+++ b/lib/psych/visitors/to_ruby.rb
@@ -271,6 +271,21 @@ module Psych
           end
           map
 
+        when /^!ruby\/marshalable:(.*)$/
+          name = $1
+          klass = resolve_class(name)
+          obj = register(o, klass.allocate)
+
+          if obj.respond_to?(:init_with)
+            init_with(obj, revive_hash({}, o), o)
+          elsif obj.respond_to?(:marshal_load)
+            marshal_data = o.children.map(&method(:accept))
+            obj.marshal_load(marshal_data)
+            obj
+          else
+            raise ArgumentError, "Cannot deserialize #{name}"
+          end
+
         else
           revive_hash(register(o, {}), o)
         end
diff --git a/lib/psych/visitors/yaml_tree.rb b/lib/psych/visitors/yaml_tree.rb
index 3f2427b..989e1f0 100644
--- a/lib/psych/visitors/yaml_tree.rb
+++ b/lib/psych/visitors/yaml_tree.rb
@@ -27,6 +27,8 @@ module Psych
 
         def key? target
           @obj_to_node.key? target.object_id
+        rescue NoMethodError
+          false
         end
 
         def id_for target
@@ -411,6 +413,18 @@ module Psych
         end
       end
 
+      def visit_BasicObject o
+        tag = Psych.dump_tags[o.class]
+        tag ||= "!ruby/marshalable:#{o.class.name}"
+
+        map = @emitter.start_mapping(nil, tag, false, Nodes::Mapping::BLOCK)
+        register(o, map)
+
+        o.marshal_dump.each(&method(:accept))
+
+        @emitter.end_mapping
+      end
+
       private
       # FIXME: Remove the index and count checks in Psych 3.0
       NULL         = "\x00"
diff --git a/test/psych/test_marshalable.rb b/test/psych/test_marshalable.rb
new file mode 100644
index 0000000..7df74ee
--- /dev/null
+++ b/test/psych/test_marshalable.rb
@@ -0,0 +1,54 @@
+require_relative 'helper'
+require 'delegate'
+
+module Psych
+  class TestMarshalable < TestCase
+    def test_objects_defining_marshal_dump_and_marshal_load_can_be_dumped
+      sd = SimpleDelegator.new(1)
+      loaded = Psych.load(Psych.dump(sd))
+
+      assert_instance_of(SimpleDelegator, loaded)
+      assert_equal(sd, loaded)
+    end
+
+    class PsychCustomMarshalable < BasicObject
+      attr_reader :foo
+
+      def initialize(foo)
+        @foo = foo
+      end
+
+      def marshal_dump
+        [foo]
+      end
+
+      def mashal_load(data)
+        @foo = data[0]
+      end
+
+      def init_with(coder)
+        @foo = coder['foo']
+      end
+
+      def encode_with(coder)
+        coder['foo'] = 2
+      end
+
+      def respond_to?(method)
+        [:marshal_dump, :marshal_load, :init_with, :encode_with].include?(method)
+      end
+
+      def class
+        PsychCustomMarshalable
+      end
+    end
+
+    def test_init_with_takes_priority_over_marshal_methods
+      obj = PsychCustomMarshalable.new(1)
+      loaded = Psych.load(Psych.dump(obj))
+
+      assert(PsychCustomMarshalable === loaded)
+      assert_equal(2, loaded.foo)
+    end
+  end
+end
-- 
cgit v1.2.3