diff --git a/lib/psych/visitors/custom_class.rb b/lib/psych/visitors/custom_class.rb new file mode 100644 index 00000000..5feacc73 --- /dev/null +++ b/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 diff --git a/lib/psych/visitors/to_ruby.rb b/lib/psych/visitors/to_ruby.rb index fd1c8e6c..856bc7d2 100644 --- a/lib/psych/visitors/to_ruby.rb +++ b/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,17 +336,25 @@ 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 @@ -349,9 +362,9 @@ def revive_hash hash, o 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 diff --git a/test/psych/visitors/test_custom_class.rb b/test/psych/visitors/test_custom_class.rb new file mode 100644 index 00000000..5ac8aeca --- /dev/null +++ b/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