Merge pull request #3986 from ksss/enum-zip

Support block yielding for Enumerable#zip

mrbgems/mruby-enum-ext/mrblib/enum.rb

@@ -742,18 +742,39 @@ module Enumerable
   ##
   #  call-seq:
   #     enum.zip(arg, ...)                  -> an_array_of_array
+  #     enum.zip(arg, ...) { |arr| block }  -> nil
   #
   #  Takes one element from <i>enum</i> and merges corresponding
   #  elements from each <i>args</i>.  This generates a sequence of
   #  <em>n</em>-element arrays, where <em>n</em> is one more than the
   #  count of arguments.  The length of the resulting sequence will be
   #  <code>enum#size</code>.  If the size of any argument is less than
-  #  <code>enum#size</code>, <code>nil</code> values are supplied.
+  #  <code>enum#size</code>, <code>nil</code> values are supplied. If
+  #  a block is given, it is invoked for each output array, otherwise
+  #  an array of arrays is returned.
+  #
+  #     a = [ 4, 5, 6 ]
+  #     b = [ 7, 8, 9 ]
+  #
+  #     a.zip(b)                 #=> [[4, 7], [5, 8], [6, 9]]
+  #     [1, 2, 3].zip(a, b)      #=> [[1, 4, 7], [2, 5, 8], [3, 6, 9]]
+  #     [1, 2].zip(a, b)         #=> [[1, 4, 7], [2, 5, 8]]
+  #     a.zip([1, 2], [8])       #=> [[4, 1, 8], [5, 2, nil], [6, nil, nil]]
+  #
+  #     c = []
+  #     a.zip(b) { |x, y| c << x + y }  #=> nil
+  #     c                               #=> [11, 13, 15]
   #
 
-  def zip(*arg)
-    ary = []
-    arg = arg.map{|a|a.to_a}
+  def zip(*arg, &block)
+    result = block ? nil : []
+    arg = arg.map do |a|
+      unless a.respond_to?(:to_a)
+        raise TypeError, "wrong argument type #{a.class} (must respond to :to_a)"
+      end
+      a.to_a
+    end
+
     i = 0
     self.each do |*val|
       a = []
@@ -763,10 +784,14 @@ module Enumerable
         a.push(arg[idx][i])
         idx += 1
       end
-      ary.push(a)
       i += 1
+      if result.nil?
+        block.call(a)
+      else
+        result.push(a)
+      end
     end
-    ary
+    result
   end
 
   ##

mrbgems/mruby-enum-ext/test/enum.rb

@@ -166,6 +166,12 @@ assert("Enumerable#zip") do
   assert_equal [[1, 4, 7], [2, 5, 8], [3, 6, 9]], [1, 2, 3].zip(a, b)
   assert_equal [[1, 4, 7], [2, 5, 8]], [1, 2].zip(a, b)
   assert_equal [[4, 1, 8], [5, 2, nil], [6, nil, nil]], a.zip([1, 2], [8])
+
+  ret = []
+  assert_equal nil, a.zip([1, 2], [8]) { |i| ret << i }
+  assert_equal [[4, 1, 8], [5, 2, nil], [6, nil, nil]], ret
+
+  assert_raise(TypeError) { [1].zip(1) }
 end
 
 assert("Enumerable#to_h") do