Special treatement for Relation#select { with block }

Author: lifo

activerecord/lib/active_record/relation/query_methods.rb

@@ -9,7 +9,7 @@ module QueryMethods
       (ActiveRecord::Relation::ASSOCIATION_METHODS + ActiveRecord::Relation::MULTI_VALUE_METHODS).each do |query_method|
         attr_accessor :"#{query_method}_values"
 
-        next if [:where, :having].include?(query_method)
+        next if [:where, :having, :select].include?(query_method)
         class_eval <<-CEVAL, __FILE__, __LINE__ + 1
           def #{query_method}(*args, &block)
             new_relation = clone
@@ -21,6 +21,19 @@ def #{query_method}(*args, &block)
         CEVAL
       end
 
+      class_eval <<-CEVAL, __FILE__, __LINE__ + 1
+        def select(*args, &block)
+          if block_given?
+            to_a.select(&block)
+          else
+            new_relation = clone
+            value = Array.wrap(args.flatten).reject {|x| x.blank? }
+            new_relation.select_values += value if value.present?
+            new_relation
+          end
+        end
+      CEVAL
+
       [:where, :having].each do |query_method|
         class_eval <<-CEVAL, __FILE__, __LINE__ + 1
           def #{query_method}(*args, &block)

activerecord/test/cases/relations_test.rb

@@ -112,6 +112,11 @@ def test_finding_with_group
     assert_equal 4, developers.map(&:salary).uniq.size
   end
 
+  def test_select_with_block
+    even_ids = Developer.scoped.select {|d| d.id % 2 == 0 }.map(&:id)
+    assert_equal [2, 4, 6, 8, 10], even_ids
+  end
+
   def test_finding_with_hash_conditions_on_joined_table
     firms = DependentFirm.joins(:account).where({:name => 'RailsCore', :accounts => { :credit_limit => 55..60 }}).to_a
     assert_equal 1, firms.size