Allow blocks for count with ActiveRecord::Relation. Document and test that sum allows blocks

2 files changed, 30 insertions, 3 deletions

diff --git a/activerecord/lib/active_record/relation/calculations.rb b/activerecord/lib/active_record/relation/calculations.rb
index 31d99f0192..a0d24f28d0 100644
--- a/activerecord/lib/active_record/relation/calculations.rb
+++ b/activerecord/lib/active_record/relation/calculations.rb
@@ -16,9 +16,16 @@ module ActiveRecord
     #
     #   Person.count(:age, distinct: true)
     #   # => counts the number of different age values
+    #
+    #   Person.where("age > 26").count { |person| gender == 'female' }
+    #   # => queries people where "age > 26" then count the loaded results filtering by gender
     def count(column_name = nil, options = {})
-      column_name, options = nil, column_name if column_name.is_a?(Hash)
-      calculate(:count, column_name, options)
+      if block_given?
+        self.to_a.count { |*block_args| yield(*block_args) }
+      else
+        column_name, options = nil, column_name if column_name.is_a?(Hash)
+        calculate(:count, column_name, options)
+      end
     end
 
     # Calculates the average value on a given column. Returns +nil+ if there's
@@ -52,9 +59,13 @@ module ActiveRecord
     # +calculate+ for examples with options.
     #
     #   Person.sum('age') # => 4562
+    #   # => returns the total sum of all people's age
+    #
+    #   Person.where('age > 100').sum { |person| person.age - 100 }
+    #   # queries people where "age > 100" then perform a sum calculation with the block returns
     def sum(*args)
       if block_given?
-        self.to_a.sum(*args) {|*block_args| yield(*block_args)}
+        self.to_a.sum(*args) { |*block_args| yield(*block_args) }
       else
         calculate(:sum, *args)
       end
diff --git a/activerecord/test/cases/calculations_test.rb b/activerecord/test/cases/calculations_test.rb
index 041f8ffb7c..a279b0e77c 100644
--- a/activerecord/test/cases/calculations_test.rb
+++ b/activerecord/test/cases/calculations_test.rb
@@ -376,6 +376,22 @@ class CalculationsTest < ActiveRecord::TestCase
         Company.where(:type => "Firm").from('companies').count(:type)
   end
 
+  def test_count_with_block_acts_as_array
+    accounts = Account.where('id > 0')
+    assert_equal Account.count, accounts.count { true }
+    assert_equal 0, accounts.count { false }
+    assert_equal Account.where('credit_limit > 50').size, accounts.count { |account| account.credit_limit > 50 }
+    assert_equal Account.count, Account.count { true }
+    assert_equal 0, Account.count { false }
+  end
+
+  def test_sum_with_block_acts_as_array
+    accounts = Account.where('id > 0')
+    assert_equal Account.sum(:credit_limit), accounts.sum { |account| account.credit_limit }
+    assert_equal Account.sum(:credit_limit) + Account.count, accounts.sum{ |account| account.credit_limit + 1 }
+    assert_equal 0, accounts.sum { |account| 0 }
+  end
+
   def test_sum_with_from_option
     assert_equal Account.sum(:credit_limit), Account.from('accounts').sum(:credit_limit)
     assert_equal Account.where("credit_limit > 50").sum(:credit_limit),