From 2f3667d0638b533756b5280e5ca2b86b17f6de93 Mon Sep 17 00:00:00 2001
From: Jason Frey <fryguy9@gmail.com>
Date: Wed, 14 Oct 2015 15:39:16 -0400
Subject: [PATCH] Add ability to pass a transformation block to element_counts

---
 .../core_ext/array/element_counts.rb           | 10 ++++++++--
 spec/core_ext/array/element_counts_spec.rb     | 18 +++++++++++++-----
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/lib/more_core_extensions/core_ext/array/element_counts.rb b/lib/more_core_extensions/core_ext/array/element_counts.rb
index e9041d3..d5f04d0 100644
--- a/lib/more_core_extensions/core_ext/array/element_counts.rb
+++ b/lib/more_core_extensions/core_ext/array/element_counts.rb
@@ -1,10 +1,16 @@
 module MoreCoreExtensions
   module ArrayElementCounts
-    # Returns a Hash of each element to the count of those elements.
+    # Returns a Hash of each element to the count of those elements.  Optionally
+    #  pass a block to count by a different criteria.
     #
     #   [1, 2, 3, 1, 3, 1].counts  # => {1 => 3, 2 => 1, 3 => 2}
+    #   %w(a aa aaa a aaa a).counts { |i| i.length }  # => {1 => 3, 2 => 1, 3 => 2}
+    #
     def element_counts
-      each_with_object(Hash.new(0)) { |i, h| h[i] += 1 }
+      each_with_object(Hash.new(0)) do |i, h|
+        key = block_given? ? yield(i) : i
+        h[key] += 1
+      end
     end
   end
 end
diff --git a/spec/core_ext/array/element_counts_spec.rb b/spec/core_ext/array/element_counts_spec.rb
index 6da5ed9..6d47ab3 100644
--- a/spec/core_ext/array/element_counts_spec.rb
+++ b/spec/core_ext/array/element_counts_spec.rb
@@ -1,8 +1,16 @@
 describe Array do
-  it "#element_counts" do
-    expect([].element_counts).to eq({})
-    expect([1].element_counts).to eq({1 => 1})
-    expect([nil].element_counts).to eq({nil => 1})
-    expect([1, 2, 3, 1, 3, 1].element_counts).to eq({1 => 3, 2 => 1, 3 => 2})
+  describe "#element_counts" do
+    it "without a block" do
+      expect([].element_counts).to eq({})
+      expect([1].element_counts).to eq({1 => 1})
+      expect([nil].element_counts).to eq({nil => 1})
+      expect([1, 2, 3, 1, 3, 1].element_counts).to eq({1 => 3, 2 => 1, 3 => 2})
+    end
+
+    it "with a block" do
+      expect([].element_counts(&:size)).to eq({})
+      expect(%w(a).element_counts(&:size)).to eq({1 => 1})
+      expect(%w(a aa aaa a aaa a).element_counts(&:size)).to eq({1 => 3, 2 => 1, 3 => 2})
+    end
   end
 end