diff --git a/mrbgems/mruby-enum-lazy/mrblib/lazy.rb b/mrbgems/mruby-enum-lazy/mrblib/lazy.rb index 7506d7fc8..140faf65b 100644 --- a/mrbgems/mruby-enum-lazy/mrblib/lazy.rb +++ b/mrbgems/mruby-enum-lazy/mrblib/lazy.rb @@ -291,11 +291,12 @@ class Enumerator # def flat_map(&block) Lazy.new(self){|yielder, val| - ary = block.call(val) - # TODO: check ary is an Array - ary.each {|x| - yielder << x - } + result = block.call(val) + if result.respond_to?(:each) + result.each {|x| yielder << x } + else + yielder << result + end } end alias collect_concat flat_map diff --git a/mrbgems/mruby-enum-lazy/test/lazy.rb b/mrbgems/mruby-enum-lazy/test/lazy.rb index 5266f6f69..145c73caf 100644 --- a/mrbgems/mruby-enum-lazy/test/lazy.rb +++ b/mrbgems/mruby-enum-lazy/test/lazy.rb @@ -46,6 +46,18 @@ assert("Enumerator::Lazy#to_enum") do assert_equal [0*1, 2*3, 4*5, 6*7], lazy_enum.map { |a| a.first * a.last }.first(4) end +assert("Enumerator::Lazy#flat_map with arrays") do + assert_equal [1, 10, 2, 20, 3, 30], [1, 2, 3].lazy.flat_map {|x| [x, x*10]}.force +end + +assert("Enumerator::Lazy#flat_map with non-enumerable") do + assert_equal [1, 2, 3], [1, 2, 3].lazy.flat_map {|x| x}.force +end + +assert("Enumerator::Lazy#flat_map with enumerable") do + assert_equal [[1, 2], [3, 4]], [1, 3].lazy.flat_map {|x| [[x, x+1]]}.force +end + assert("Enumerator::Lazy#grep_v") do lazy_grep_v = (0..).lazy.grep_v(2..4) assert_kind_of Enumerator::Lazy, lazy_grep_v