mirror of
https://github.com/mruby/mruby
synced 2026-06-08 16:11:16 +00:00
Fix Lazy#flat_map to handle non-enumerable block return values
When the block passed to Lazy#flat_map returns a non-enumerable value (e.g. an Integer), mruby raised NoMethodError because it unconditionally called #each on the result. CRuby yields non-enumerable values directly. Use respond_to?(:each) to match CRuby behavior: iterate enumerable results, yield non-enumerable results as-is.
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user