[ruby/psych] fixed bugs from testing

https://github.com/ruby/psych/commit/e954f96639

1 files changed, 11 insertions, 2 deletions

diff --git a/ext/psych/lib/psych.rb b/ext/psych/lib/psych.rb
index 64637b97f7..1c9dea5bba 100644
--- a/ext/psych/lib/psych.rb
+++ b/ext/psych/lib/psych.rb
@@ -668,8 +668,17 @@ module Psych
   #   list # => ['foo', 'bar']
   #
   def self.safe_load_stream yaml, filename: nil, permitted_classes: [], aliases: false
-    parse_stream(yaml, filename: filename).children.map do |child|
-      safe_load(child.to_yaml, permitted_classes, aliases: aliases)
+    documents = parse_stream(yaml, filename: filename).children.map do |child|
+      stream = Psych::Nodes::Stream.new
+      stream.children << child
+      safe_load stream.to_yaml, permitted_classes: permitted_classes, aliases: aliases
+    end
+
+    if block_given?
+      documents.each { |doc| yield doc }
+      nil
+    else
+      documents
     end
   end