Add block setting of attributes to singular associations

5 files changed, 64 insertions, 11 deletions

diff --git a/activerecord/CHANGELOG b/activerecord/CHANGELOG
index 2e144745cd..502a7e43de 100644
--- a/activerecord/CHANGELOG
+++ b/activerecord/CHANGELOG
@@ -1,5 +1,15 @@
 *Rails 3.1.0 (unreleased)*
 
+* Add block setting of attributes to singular associations:
+
+    class User < ActiveRecord::Base
+      has_one :account
+    end
+
+    user.build_account{ |a| a.credit_limit => 100.0 }
+
+  The block is called after the instance has been initialized. [Andrew White]
+
 * Add ActiveRecord::Base.attribute_names to return a list of attribute names. This will return an empty array if the model is abstract or table does not exists. [Prem Sichanugrist]
 
 * CSV Fixtures are deprecated and support will be removed in Rails 3.2.0
diff --git a/activerecord/lib/active_record/associations/builder/singular_association.rb b/activerecord/lib/active_record/associations/builder/singular_association.rb
index 62d48d3a2c..638a2ec72a 100644
--- a/activerecord/lib/active_record/associations/builder/singular_association.rb
+++ b/activerecord/lib/active_record/associations/builder/singular_association.rb
@@ -29,16 +29,16 @@ module ActiveRecord::Associations::Builder
       def define_constructors
         name = self.name
 
-        model.redefine_method("build_#{name}") do |*params|
-          association(name).build(*params)
+        model.redefine_method("build_#{name}") do |*params, &block|
+          association(name).build(*params, &block)
         end
 
-        model.redefine_method("create_#{name}") do |*params|
-          association(name).create(*params)
+        model.redefine_method("create_#{name}") do |*params, &block|
+          association(name).create(*params, &block)
         end
 
-        model.redefine_method("create_#{name}!") do |*params|
-          association(name).create!(*params)
+        model.redefine_method("create_#{name}!") do |*params, &block|
+          association(name).create!(*params, &block)
         end
       end
   end
diff --git a/activerecord/lib/active_record/associations/singular_association.rb b/activerecord/lib/active_record/associations/singular_association.rb
index 44dbe984d4..68fe0bde76 100644
--- a/activerecord/lib/active_record/associations/singular_association.rb
+++ b/activerecord/lib/active_record/associations/singular_association.rb
@@ -17,17 +17,18 @@ module ActiveRecord
         replace(record)
       end
 
-      def create(attributes = {}, options = {})
-        build(attributes, options).tap { |record| record.save }
+      def create(attributes = {}, options = {}, &block)
+        build(attributes, options, &block).tap { |record| record.save }
       end
 
-      def create!(attributes = {}, options = {})
-        build(attributes, options).tap { |record| record.save! }
+      def create!(attributes = {}, options = {}, &block)
+        build(attributes, options, &block).tap { |record| record.save! }
       end
 
-      def build(attributes = {}, options = {})
+      def build(attributes = {}, options = {}, &block)
         record = reflection.build_association(attributes, options)
         record.assign_attributes(create_scope.except(*record.changed), :without_protection => true)
+        yield(record) if block_given?
         set_new_record(record)
         record
       end
diff --git a/activerecord/test/cases/associations/belongs_to_associations_test.rb b/activerecord/test/cases/associations/belongs_to_associations_test.rb
index ddcc36c841..b993bf6e90 100644
--- a/activerecord/test/cases/associations/belongs_to_associations_test.rb
+++ b/activerecord/test/cases/associations/belongs_to_associations_test.rb
@@ -626,4 +626,25 @@ class BelongsToAssociationsTest < ActiveRecord::TestCase
 
     assert_equal "Bob", firm.name
   end
+
+  def test_build_with_block
+    client = Client.create(:name => 'Client Company')
+
+    firm = client.build_firm{ |f| f.name = 'Agency Company' }
+    assert_equal 'Agency Company', firm.name
+  end
+
+  def test_create_with_block
+    client = Client.create(:name => 'Client Company')
+
+    firm = client.create_firm{ |f| f.name = 'Agency Company' }
+    assert_equal 'Agency Company', firm.name
+  end
+
+  def test_create_bang_with_block
+    client = Client.create(:name => 'Client Company')
+
+    firm = client.create_firm!{ |f| f.name = 'Agency Company' }
+    assert_equal 'Agency Company', firm.name
+  end
 end
diff --git a/activerecord/test/cases/associations/has_one_associations_test.rb b/activerecord/test/cases/associations/has_one_associations_test.rb
index 38e622f90f..f3bf5baa95 100644
--- a/activerecord/test/cases/associations/has_one_associations_test.rb
+++ b/activerecord/test/cases/associations/has_one_associations_test.rb
@@ -426,4 +426,25 @@ class HasOneAssociationsTest < ActiveRecord::TestCase
     bulb = car.build_bulb({ :bulb_type => :custom }, :as => :admin)
     assert_equal CustomBulb, bulb.class
   end
+
+  def test_build_with_block
+    car = Car.create(:name => 'honda')
+
+    bulb = car.build_bulb{ |b| b.color = 'Red' }
+    assert_equal 'RED!', bulb.color
+  end
+
+  def test_create_with_block
+    car = Car.create(:name => 'honda')
+
+    bulb = car.create_bulb{ |b| b.color = 'Red' }
+    assert_equal 'RED!', bulb.color
+  end
+
+  def test_create_bang_with_block
+    car = Car.create(:name => 'honda')
+
+    bulb = car.create_bulb!{ |b| b.color = 'Red' }
+    assert_equal 'RED!', bulb.color
+  end
 end