diff options
| -rw-r--r-- | mrbgems/mruby-enum-ext/mrblib/enum.rb | 30 | ||||
| -rw-r--r-- | mrbgems/mruby-enum-ext/test/enum.rb | 4 |
2 files changed, 29 insertions, 5 deletions
diff --git a/mrbgems/mruby-enum-ext/mrblib/enum.rb b/mrbgems/mruby-enum-ext/mrblib/enum.rb index 6a978b1e3..96249db28 100644 --- a/mrbgems/mruby-enum-ext/mrblib/enum.rb +++ b/mrbgems/mruby-enum-ext/mrblib/enum.rb @@ -742,18 +742,34 @@ 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 = [] + def zip(*arg, &block) + result = block ? nil : [] arg = arg.map{|a|a.to_a} + i = 0 self.each do |*val| a = [] @@ -763,10 +779,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 ## diff --git a/mrbgems/mruby-enum-ext/test/enum.rb b/mrbgems/mruby-enum-ext/test/enum.rb index 7b9989861..ee4d8b844 100644 --- a/mrbgems/mruby-enum-ext/test/enum.rb +++ b/mrbgems/mruby-enum-ext/test/enum.rb @@ -166,6 +166,10 @@ 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 end assert("Enumerable#to_h") do |