Handle non array return values in Enumerable#flat_map (#3427)

crystal-lang · Oct 18, 2016 · 8d326cc · 8d326cc
1 parent c020630
commit 8d326cc
Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/spec/std/enumerable_spec.cr b/spec/std/enumerable_spec.cr
@@ -237,6 +237,14 @@ describe "Enumerable" do
     it "does example 2" do
       [[1, 2], [3, 4]].flat_map { |e| e + [100] }.should eq([1, 2, 100, 3, 4, 100])
     end
+
+    it "does example 3" do
+      [[1, 2, 3], 4, 5].flat_map { |e| e }.should eq([1, 2, 3, 4, 5])
+    end
+
+    it "does example 4" do
+      [{1 => 2}, {3 => 4}].flat_map { |e| e }.should eq([{1 => 2}, {3 => 4}])
+    end
   end
 
   describe "grep" do

diff --git a/src/enumerable.cr b/src/enumerable.cr
@@ -250,9 +250,16 @@ module Enumerable(T)
   #     ["Alice", "Bob"].flat_map do |user|
   #       user.chars
   #     end  #=> ['A', 'l', 'i', 'c', 'e', 'B', 'o', 'b']
-  def flat_map(&block : T -> Array(U)) forall U
+  def flat_map(&block : T -> Array(U) | U) forall U
     ary = [] of U
-    each { |e| ary.concat(yield e) }
+    each do |e|
+      v = yield e
+      if v.is_a?(Array)
+        ary.concat(v)
+      else
+        ary.push(v)
+      end
+    end
     ary
   end