Implement freeze option for Pysch.load

Author: byroot

lib/psych/visitors/to_ruby.rb

@@ -32,16 +32,18 @@ def initialize ss, class_loader, symbolize_names: false, freeze: false
 
       def accept target
         result = super
-        return result if @domain_types.empty? || !target.tag
 
-        key = target.tag.sub(/^[!\/]*/, '').sub(/(,\d+)\//, '\1:')
-        key = "tag:#{key}" unless key =~ /^(?:tag:|x-private)/
+        unless @domain_types.empty? || !target.tag
+          key = target.tag.sub(/^[!\/]*/, '').sub(/(,\d+)\//, '\1:')
+          key = "tag:#{key}" unless key =~ /^(?:tag:|x-private)/
 
-        if @domain_types.key? key
-          value, block = @domain_types[key]
-          return block.call value, result
+          if @domain_types.key? key
+            value, block = @domain_types[key]
+            result = block.call value, result
+          end
         end
 
+        result = deduplicate(result).freeze if @freeze
         result
       end
 
@@ -341,7 +343,7 @@ def revive_hash hash, o
           key = accept(k)
           if @symbolize_names
             key = key.to_sym
-          else
+          elsif !@freeze
             key = deduplicate(key)
           end
           val = accept(v)
@@ -378,6 +380,8 @@ def revive_hash hash, o
       if RUBY_VERSION < '2.7'
         def deduplicate key
           if key.is_a?(String)
+            # It is important to untaint the string, otherwise it won't
+            # be deduplicated into an fstring, but simply frozen.
             -(key.untaint)
           else
             key

test/psych/test_psych.rb

@@ -192,6 +192,22 @@ def test_domain_types
     assert_equal({ 'hello' => 'world' }, got)
   end
 
+  def test_load_freeze
+    data = Psych.load("--- {foo: ['a']}", freeze: true)
+    assert_predicate data, :frozen?
+    assert_predicate data['foo'], :frozen?
+    assert_predicate data['foo'].first, :frozen?
+  end
+
+  def test_load_freeze_deduplication
+    unless String.method_defined?(:-@) && (-("a" * 20)).equal?((-("a" * 20)))
+      skip "This Ruby implementation doesn't support string deduplication"
+    end
+
+    data = Psych.load("--- ['a']", freeze: true)
+    assert_same 'a', data.first
+  end
+
   def test_load_default_fallback
     assert_equal false, Psych.load("")
   end