Merge pull request #32523 from kaspth/enumerable-index-with-extension

Add Enumerable#index_with.

4 files changed, 70 insertions, 1 deletions

diff --git a/activesupport/CHANGELOG.md b/activesupport/CHANGELOG.md
index 62c0f612a8..a384fc3bd2 100644
--- a/activesupport/CHANGELOG.md
+++ b/activesupport/CHANGELOG.md
@@ -1,3 +1,18 @@
+*   Add `index_with` to Enumerable.
+
+    Allows creating a hash from an enumerable with the value from a passed block
+    or a default argument.
+
+        %i( title body ).index_with { |attr| post.public_send(attr) }
+        # => { title: "hey", body: "what's up?" }
+
+        %i( title body ).index_with(nil)
+        # => { title: nil, body: nil }
+
+    Closely linked with its brethen `index_by`.
+
+    *Kasper Timm Hansen*
+
 *   Fix bug where `ActiveSupport::Timezone.all` would fail when tzinfo data for
     any timezone defined in `ActiveSupport::TimeZone::MAPPING` is missing.
 
diff --git a/activesupport/lib/active_support/core_ext/enumerable.rb b/activesupport/lib/active_support/core_ext/enumerable.rb
index edde4f46b9..7713c52cb1 100644
--- a/activesupport/lib/active_support/core_ext/enumerable.rb
+++ b/activesupport/lib/active_support/core_ext/enumerable.rb
@@ -1,6 +1,9 @@
 # frozen_string_literal: true
 
 module Enumerable
+  INDEX_WITH_DEFAULT = Object.new
+  private_constant :INDEX_WITH_DEFAULT
+
   # Enumerable#sum was added in Ruby 2.4, but it only works with Numeric elements
   # when we omit an identity.
 
@@ -37,10 +40,11 @@ module Enumerable
     end
   end
 
-  # Convert an enumerable to a hash.
+  # Convert an enumerable to a hash keying it by the block return value.
   #
   #   people.index_by(&:login)
   #   # => { "nextangle" => <Person ...>, "chade-" => <Person ...>, ...}
+  #
   #   people.index_by { |person| "#{person.first_name} #{person.last_name}" }
   #   # => { "Chade- Fowlersburg-e" => <Person ...>, "David Heinemeier Hansson" => <Person ...>, ...}
   def index_by
@@ -53,6 +57,26 @@ module Enumerable
     end
   end
 
+  # Convert an enumerable to a hash keying it with the enumerable items and with the values returned in the block.
+  #
+  #   post = Post.new(title: "hey there", body: "what's up?")
+  #
+  #   %i( title body ).index_with { |attr_name| post.public_send(attr_name) }
+  #   # => { title: "hey there", body: "what's up?" }
+  def index_with(default = INDEX_WITH_DEFAULT)
+    if block_given?
+      result = {}
+      each { |elem| result[elem] = yield(elem) }
+      result
+    elsif default != INDEX_WITH_DEFAULT
+      result = {}
+      each { |elem| result[elem] = default }
+      result
+    else
+      to_enum(:index_with) { size if respond_to?(:size) }
+    end
+  end
+
   # Returns +true+ if the enumerable has more than 1 element. Functionally
   # equivalent to <tt>enum.to_a.size > 1</tt>. Can be called with a block too,
   # much like any?, so <tt>people.many? { |p| p.age > 26 }</tt> returns +true+
diff --git a/activesupport/test/core_ext/enumerable_test.rb b/activesupport/test/core_ext/enumerable_test.rb
index 8d71320931..b63464a36a 100644
--- a/activesupport/test/core_ext/enumerable_test.rb
+++ b/activesupport/test/core_ext/enumerable_test.rb
@@ -179,6 +179,21 @@ class EnumerableTests < ActiveSupport::TestCase
                  payments.index_by.each(&:price))
   end
 
+  def test_index_with
+    payments = GenericEnumerable.new([ Payment.new(5), Payment.new(15), Payment.new(10) ])
+
+    assert_equal({ Payment.new(5) => 5, Payment.new(15) => 15, Payment.new(10) => 10 }, payments.index_with(&:price))
+
+    assert_equal({ title: nil, body: nil }, %i( title body ).index_with(nil))
+    assert_equal({ title: [], body: [] }, %i( title body ).index_with([]))
+    assert_equal({ title: {}, body: {} }, %i( title body ).index_with({}))
+
+    assert_equal Enumerator, payments.index_with.class
+    assert_nil payments.index_with.size
+    assert_equal 42, (1..42).index_with.size
+    assert_equal({ Payment.new(5) => 5, Payment.new(15) => 15, Payment.new(10) => 10 }, payments.index_with.each(&:price))
+  end
+
   def test_many
     assert_equal false, GenericEnumerable.new([]).many?
     assert_equal false, GenericEnumerable.new([ 1 ]).many?
diff --git a/guides/source/active_support_core_extensions.md b/guides/source/active_support_core_extensions.md
index ae2e1faf14..cde217e1e4 100644
--- a/guides/source/active_support_core_extensions.md
+++ b/guides/source/active_support_core_extensions.md
@@ -2039,6 +2039,21 @@ WARNING. Keys should normally be unique. If the block returns the same value for
 
 NOTE: Defined in `active_support/core_ext/enumerable.rb`.
 
+### `index_with`
+
+The method `index_with` generates a hash with the elements of an enumerable as keys. The value
+is either a passed default or returned in a block.
+
+```ruby
+%i( title body created_at ).index_with { |attr_name| public_send(attr_name) }
+# => { title: "hey", body: "what's up?", … }
+
+WEEKDAYS.index_with([ Interval.all_day ])
+# => { monday: [ 0, 1440 ], … }
+```
+
+NOTE: Defined in `active_support/core_ext/enumerable.rb`.
+
 ### `many?`
 
 The method `many?` is shorthand for `collection.size > 1`: