Make `#to_h` to take a block; [ruby-core:89088]

2 files changed, 17 insertions, 6 deletions

diff --git a/mrbgems/mruby-enum-ext/mrblib/enum.rb b/mrbgems/mruby-enum-ext/mrblib/enum.rb
index a840ade3b..6cbacdf9e 100644
--- a/mrbgems/mruby-enum-ext/mrblib/enum.rb
+++ b/mrbgems/mruby-enum-ext/mrblib/enum.rb
@@ -803,13 +803,22 @@ module Enumerable
   #       # => {:hello => 0, :world => 1}
   #
 
-  def to_h
+  def to_h(&blk)
     h = {}
-    self.each do |*v|
-      v = v.__svalue
-      raise TypeError, "wrong element type #{v.class} (expected Array)" unless v.is_a? Array
-      raise ArgumentError, "element has wrong array length (expected 2, was #{v.size})" if v.size != 2
-      h[v[0]] = v[1]
+    if blk
+      self.each do |v|
+        v = blk.call(v)
+        raise TypeError, "wrong element type #{v.class} (expected Array)" unless v.is_a? Array
+        raise ArgumentError, "element has wrong array length (expected 2, was #{v.size})" if v.size != 2
+        h[v[0]] = v[1]
+      end
+    else
+      self.each do |*v|
+        v = v.__svalue
+        raise TypeError, "wrong element type #{v.class} (expected Array)" unless v.is_a? Array
+        raise ArgumentError, "element has wrong array length (expected 2, was #{v.size})" if v.size != 2
+        h[v[0]] = v[1]
+      end
     end
     h
   end
diff --git a/mrbgems/mruby-enum-ext/test/enum.rb b/mrbgems/mruby-enum-ext/test/enum.rb
index 46ed5f0f9..b2d7297a4 100644
--- a/mrbgems/mruby-enum-ext/test/enum.rb
+++ b/mrbgems/mruby-enum-ext/test/enum.rb
@@ -188,4 +188,6 @@ assert("Enumerable#to_h") do
   assert_equal h0, h
   # mruby-enum-ext also provides nil.to_h
   assert_equal Hash.new, nil.to_h
+
+  assert_equal({1=>4,3=>8}, c.new.to_h{|k,v|[k,v*2]})
 end