feat(visitor): allow to create custom object (instead of Hash) according to context

lib/psych/visitors/custom_class.rb

@@ -0,0 +1,122 @@
+# frozen_string_literal: false
+
+require 'psych/visitors/to_ruby'
+
+
+module Psych
+  module Visitors
+
+    ##
+    ## Visitor class to generate custom object instead of Hash.
+    ##
+    ## Example1:
+    ##
+    ##     ## define custom classes
+    ##     Team   = Struct.new('Team',   'name', 'members')
+    ##     Member = Struct.new('Member', 'name', 'gender')
+    ##     ## create visitor object
+    ##     require 'psych'
+    ##     require 'psych/visitors/custom_class'
+    ##     classmap = {
+    ##       "teams"   => Team,
+    ##       "members" => Member,
+    ##     }
+    ##     visitor = Psych::Visitors::CustomClassVisitor.create(classmap)
+    ##     ## sample YAML document
+    ##     input = <<-'END'
+    ##     teams:
+    ##       - name: SOS Brigade
+    ##         members:
+    ##           - {name: Haruhi, gender: F}
+    ##           - {name: Kyon,   gender: M}
+    ##           - {name: Mikuru, gender: F}
+    ##           - {name: Itsuki, gender: M}
+    ##           - {name: Yuki,   gender: F}
+    ##     END
+    ##     ## parse YAML document with custom classes
+    ##     tree = Psych.parse(input)
+    ##     ydoc = visitor.accept(tree)
+    ##     p ydoc['teams'][0].class                #=> Struct::Team
+    ##     p ydoc['teams'][0]['members'][0].class  #=> Struct::Member
+    ##     team = ydoc['teams'][0]
+    ##     p team.name                #=> "SOS Brigade"
+    ##     p team.members[0].name     #=> "Haruhi"
+    ##     p team.members[0].gender   #=> "F"
+    ##
+    ## Example2:
+    ##
+    ##     ## allows `hash.foo` instead of `hash["foo"]`
+    ##     class MagicHash < Hash
+    ##       def method_missing(method, *args)
+    ##         return super unless args.empty?
+    ##         return self[method.to_s]
+    ##       end
+    ##     end
+    ##     ## create visitor with custom hash class
+    ##     require 'psych'
+    ##     require 'psych/visitors/custom_class'
+    ##     classmap = {'*' => MagicHash}
+    ##     visitor = Psych::Visitors::CustomClassVisitor.create(classmap)
+    ##     ## sample YAML document
+    ##     input = <<-'END'
+    ##     teams:
+    ##       - name: SOS Brigade
+    ##         members:
+    ##           - {name: Haruhi, gender: F}
+    ##           - {name: Kyon,   gender: M}
+    ##           - {name: Mikuru, gender: F}
+    ##           - {name: Itsuki, gender: M}
+    ##           - {name: Yuki,   gender: F}
+    ##     END
+    ##     ## parse YAML document with custom hash class
+    ##     tree = Psych.parse(input)
+    ##     ydoc = visitor.accept(tree)
+    ##     p ydoc.class                            #=> MagicHash
+    ##     p ydoc['teams'][0].class                #=> MagicHash
+    ##     p ydoc['teams'][0]['members'][0].class  #=> MagicHash
+    ##     p ydoc.teams[0].members[0].name    #=> "Haruhi"
+    ##     p ydoc.teams[0].members[0].gender  #=> "F"
+    ##
+
+    class CustomClassVisitor < ToRuby
+
+      def self.create(classmap={})
+        visitor = super()
+        visitor.instance_variable_set('@classmap', classmap)
+        visitor
+      end
+
+      attr_reader :classmap   # key: string, value: class object
+
+      def initialize(*args)
+        super
+        @key_path = []     # ex: [] -> ['tables'] -> ['tables', 'columns']
+      end
+
+      private
+
+      def accept_key(k)    # push keys
+        key = super k
+        @key_path << key
+        key
+      end
+
+      def accept_value(v)  # pop keys
+        value = super v
+        @key_path.pop()
+        value
+      end
+
+      def empty_mapping(o)  # generate custom object (or Hash object)
+        klass = @classmap[@key_path.last] || @classmap['*']
+        klass ? klass.new : super
+      end
+
+      def merge_mapping(hash, val)   # for '<<' (merge)
+        val.each {|k, v| hash[k] = v }
+      end
+
+    end
+
+  end
+end

lib/psych/visitors/to_ruby.rb

@@ -142,7 +142,7 @@ def visit_Psych_Nodes_Sequence o
         when '!omap', 'tag:yaml.org,2002:omap'
           map = register(o, Psych::Omap.new)
           o.children.each { |a|
-            map[accept(a.children.first)] = accept a.children.last
+            map[accept_key(a.children.first)] = accept_value a.children.last
           }
           map
         when /^!(?:seq|ruby\/array):(.*)$/
@@ -159,7 +159,7 @@ def visit_Psych_Nodes_Mapping o
         if Psych.load_tags[o.tag]
           return revive(resolve_class(Psych.load_tags[o.tag]), o)
         end
-        return revive_hash(register(o, {}), o) unless o.tag
+        return revive_hash(register(o, empty_mapping(o)), o) unless o.tag
 
         case o.tag
         when /^!ruby\/struct:?(.*)?$/
@@ -171,8 +171,8 @@ def visit_Psych_Nodes_Mapping o
             members = {}
             struct_members = s.members.map { |x| class_loader.symbolize x }
             o.children.each_slice(2) do |k,v|
-              member = accept(k)
-              value  = accept(v)
+              member = accept_key(k)
+              value  = accept_value(v)
               if struct_members.include?(class_loader.symbolize(member))
                 s.send("#{member}=", value)
               else
@@ -215,8 +215,8 @@ def visit_Psych_Nodes_Mapping o
           string  = nil
 
           o.children.each_slice(2) do |k,v|
-            key   = accept k
-            value = accept v
+            key   = accept_key k
+            value = accept_value v
 
             if key == 'str'
               if klass
@@ -258,7 +258,7 @@ def visit_Psych_Nodes_Mapping o
           set = class_loader.psych_set.new
           @st[o.anchor] = set if o.anchor
           o.children.each_slice(2) do |k,v|
-            set[accept(k)] = accept(v)
+            set[accept_key(k)] = accept_value(v)
           end
           set
 
@@ -271,7 +271,7 @@ def visit_Psych_Nodes_Mapping o
               revive_hash hash, value
             when 'ivars'
               value.children.each_slice(2) do |k,v|
-                hash.instance_variable_set accept(k), accept(v)
+                hash.instance_variable_set accept_key(k), accept_value(v)
               end
             end
           end
@@ -283,7 +283,7 @@ def visit_Psych_Nodes_Mapping o
         when '!omap', 'tag:yaml.org,2002:omap'
           map = register(o, class_loader.psych_omap.new)
           o.children.each_slice(2) do |l,r|
-            map[accept(l)] = accept r
+            map[accept_key(l)] = accept_value r
           end
           map
 
@@ -303,7 +303,7 @@ def visit_Psych_Nodes_Mapping o
           end
 
         else
-          revive_hash(register(o, {}), o)
+          revive_hash(register(o, empty_mapping(o)), o)
         end
       end
 
@@ -320,6 +320,11 @@ def visit_Psych_Nodes_Alias o
       end
 
       private
+
+      def empty_mapping o
+        return {}
+      end
+
       def register node, object
         @st[node.anchor] = object if node.anchor
         object
@@ -331,27 +336,35 @@ def register_empty object
         list
       end
 
+      def accept_key k
+        accept(k)
+      end
+
+      def accept_value v
+        accept(v)
+      end
+
       SHOVEL = '<<'
       def revive_hash hash, o
         o.children.each_slice(2) { |k,v|
-          key = accept(k)
-          val = accept(v)
+          key = accept_key(k)
+          val = accept_value(v)
 
           if key == SHOVEL && k.tag != "tag:yaml.org,2002:str"
             case v
             when Nodes::Alias, Nodes::Mapping
               begin
-                hash.merge! val
+                merge_mapping(hash, val)
               rescue TypeError
                 hash[key] = val
               end
             when Nodes::Sequence
               begin
                 h = {}
                 val.reverse_each do |value|
-                  h.merge! value
+                  merge_mapping(h, value)
                 end
-                hash.merge! h
+                merge_mapping(hash, h)
               rescue TypeError
                 hash[key] = val
               end
@@ -366,6 +379,10 @@ def revive_hash hash, o
         hash
       end
 
+      def merge_mapping hash, val
+        hash.merge! val
+      end
+
       def merge_key hash, key, val
       end
 

test/psych/visitors/test_custom_class.rb

@@ -0,0 +1,98 @@
+# coding: US-ASCII
+# frozen_string_literal: false
+require 'psych/helper'
+require 'psych/visitors/custom_class'
+
+module Psych
+  module Visitors
+    class TestCustomClass < TestCase
+
+      INPUT_STRING = <<-'END'
+        teams:
+          - name: SOS Brigade
+            members:
+              - {name: Haruhi, gender: F}
+              - {name: Kyon,   gender: M}
+              - {name: Mikuru, gender: F}
+              - {name: Itsuki, gender: M}
+              - {name: Yuki,   gender: F}
+      END
+
+      def test_custom_classes
+        classmap = {
+          "teams"   => Struct.new('Team', 'name', 'members'),
+          "members" => Struct.new('Member', 'name', 'gender'),
+        }
+        #
+        visitor = Psych::Visitors::CustomClassVisitor.create(classmap)
+        tree = Psych.parse(INPUT_STRING)
+        ydoc = visitor.accept(tree)
+        #
+        assert_kind_of Hash,       ydoc
+        assert_kind_of classmap["teams"],   ydoc['teams'][0]
+        assert_kind_of classmap["members"], ydoc['teams'][0]['members'][0]
+        #
+        team = ydoc['teams'][0]
+        assert_equal 'SOS Brigade', team.name
+        assert_equal 'Haruhi',      team.members[0].name
+        assert_equal 'F',           team.members[0].gender
+      end
+
+      def test_default_class
+        magic_hash_cls = Class.new(Hash) do
+          def method_missing(method, *args)
+            return super unless args.empty?
+            return self[method.to_s]
+          end
+        end
+        classmap = {'*' => magic_hash_cls}
+        #
+        visitor = Psych::Visitors::CustomClassVisitor.create(classmap)
+        tree = Psych.parse(INPUT_STRING)
+        ydoc = visitor.accept(tree)
+        #
+        assert_kind_of magic_hash_cls, ydoc
+        assert_kind_of magic_hash_cls, ydoc['teams'][0]
+        assert_kind_of magic_hash_cls, ydoc['teams'][0]['members'][0]
+        #
+        team = ydoc['teams'][0]
+        assert_equal "SOS Brigade", team.name
+        assert_equal "Haruhi",      team.members[0].name
+        assert_equal "F",           team.members[0].gender
+      end
+
+      def test_merge_mapping
+        input = <<-END
+        column-defaults:
+          - &id
+            name  : id
+            type  : int
+            pkey  : true
+        tables:
+          - name  : admin_users
+            columns:
+              - <<: *id
+                name:  user_id
+        END
+        #
+        classmap = {
+          "tables"  => Struct.new('Table', 'name', 'columns'),
+          "columns" => Struct.new('Column', 'name', 'type', 'pkey', 'required'),
+        }
+        #
+        visitor = Psych::Visitors::CustomClassVisitor.create(classmap)
+        tree = Psych.parse(input)
+        ydoc = visitor.accept(tree)
+        #
+        assert_kind_of classmap["tables"],  ydoc['tables'][0]
+        assert_kind_of classmap["columns"], ydoc['tables'][0]['columns'][0]
+        #
+        table = ydoc['tables'][0]
+        assert_equal "int",     table.columns[0].type   # merged
+        assert_equal true,      table.columns[0].pkey   # merged
+        assert_equal "user_id", table.columns[0].name   # ovrerwritten
+      end
+
+    end
+  end
+end