diff --git a/spec/std/enumerable_spec.cr b/spec/std/enumerable_spec.cr index eddb921e9f7c16e0b15b8ed41e8e9aaaa1328ad6..4a312786545576196027ef414782e4280ec190ba 100755 --- a/spec/std/enumerable_spec.cr +++ b/spec/std/enumerable_spec.cr @@ -83,4 +83,8 @@ describe "Enumerable" do assert { (1..3).sum { |x| x * 2 }.should eq(12) } assert { (1..3).sum(1.5) { |x| x * 2 }.should eq(13.5) } end + + describe "compact map" do + assert { [1, nil, 2, nil, 3].compact_map { |x| x.try &.+(1) }.should eq([2, 3, 4]) } + end end diff --git a/src/array.cr b/src/array.cr index c548686cdbb4c30d54d16114a9c63dc197dbaa8d..4fc6a6b038b8ef078104217a670ffac7ad2544ca 100644 --- a/src/array.cr +++ b/src/array.cr @@ -274,7 +274,7 @@ class Array(T) end def compact - select { |x| !x.nil? } + compact_map { |x| x } end def compact! diff --git a/src/enumerable.cr b/src/enumerable.cr index fdbeb060d6510aeb5a643984b67e6857de0f8d8b..4d3976804b09f73944a63430d09aa643badc0937 100644 --- a/src/enumerable.cr +++ b/src/enumerable.cr @@ -65,6 +65,15 @@ module Enumerable(T) ary end + def compact_map(&block : T -> Nil | U) + ary = [] of U + each do |e| + v = yield e + ary << v if v + end + ary + end + def select(&block : T ->) ary = [] of T each { |e| ary << e if yield e }