Accept blocks passed to new/build/create/create! on documents and ass…

…ociations. Closes #352.

lib/mongo_mapper/plugins/associations/belongs_to_proxy.rb

@@ -18,16 +18,16 @@ def replace(doc)
           @target
         end
 
-        def build(attrs={})
-          instantiate_target(:new, attrs)
+        def build(attrs={}, &block)
+          instantiate_target(:new, attrs, &block)
         end
 
-        def create(attrs={})
-          instantiate_target(:create, attrs)
+        def create(attrs={}, &block)
+          instantiate_target(:create, attrs, &block)
         end
 
-        def create!(attrs={})
-          instantiate_target(:create!, attrs)
+        def create!(attrs={}, &block)
+          instantiate_target(:create!, attrs, &block)
         end
 
         def save_to_collection(options={})
@@ -40,8 +40,8 @@ def find_target
             klass.find_by_id(proxy_owner[association.foreign_key])
           end
 
-          def instantiate_target(instantiator, attrs={})
-            @target = klass.send(instantiator, attrs)
+          def instantiate_target(instantiator, attrs={}, &block)
+            @target = klass.send(instantiator, attrs, &block)
             proxy_owner[association.foreign_key] = @target.id
             loaded
             @target

lib/mongo_mapper/plugins/associations/many_documents_proxy.rb

@@ -9,15 +9,15 @@ class ManyDocumentsProxy < Collection
 
         def replace(docs)
           load_target
-          
+
           (target - docs).each do |t|
             case options[:dependent]
               when :destroy    then t.destroy
               when :delete_all then t.delete
               else t.update_attributes(self.foreign_key => nil)
             end
           end
-          
+
           docs.each { |doc| prepare(doc).save }
           reset
         end
@@ -33,21 +33,26 @@ def <<(*docs)
         def build(attrs={})
           doc = klass.new(attrs)
           apply_scope(doc)
+          yield doc if block_given?
           @target ||= [] unless loaded?
           @target << doc
           doc
         end
 
         def create(attrs={})
           doc = klass.new(attrs)
-          apply_scope(doc).save
+          apply_scope(doc)
+          yield doc if block_given?
+          doc.save
           reset
           doc
         end
 
         def create!(attrs={})
           doc = klass.new(attrs)
-          apply_scope(doc).save!
+          apply_scope(doc)
+          yield doc if block_given?
+          doc.save!
           reset
           doc
         end

lib/mongo_mapper/plugins/associations/one_proxy.rb

@@ -3,16 +3,16 @@ module MongoMapper
   module Plugins
     module Associations
       class OneProxy < Proxy
-        def build(attrs={})
-          instantiate_target(:new, attrs)
+        def build(attrs={}, &block)
+          instantiate_target(:new, attrs, &block)
         end
 
-        def create(attrs={})
-          instantiate_target(:create, attrs)
+        def create(attrs={}, &block)
+          instantiate_target(:create, attrs, &block)
         end
 
-        def create!(attrs={})
-          instantiate_target(:create!, attrs)
+        def create!(attrs={}, &block)
+          instantiate_target(:create!, attrs, &block)
         end
 
         def replace(doc)
@@ -29,7 +29,7 @@ def replace(doc)
               end
             end
           end
-          
+
           unless doc.nil?
             proxy_owner.save unless proxy_owner.persisted?
             doc = klass.new(doc) unless doc.is_a?(klass)
@@ -40,17 +40,17 @@ def replace(doc)
           loaded
           @target = doc
         end
-        
+
         def destroy
           target.destroy
           reset
         end
-        
+
         def delete
           target.delete
           reset
         end
-        
+
         def nullify
           nullify_scope(target)
           target.save
@@ -62,8 +62,8 @@ def find_target
             target_class.first(association.query_options.merge(criteria))
           end
 
-          def instantiate_target(instantiator, attrs={})
-            @target = target_class.send(instantiator, attrs.update(criteria))
+          def instantiate_target(instantiator, attrs={}, &block)
+            @target = target_class.send(instantiator, attrs.update(criteria), &block)
             loaded
             @target
           end

lib/mongo_mapper/plugins/associations/single_association.rb

@@ -27,16 +27,16 @@ def #{name}?
               get_proxy(associations[#{name.inspect}]).present?
             end
 
-            def build_#{name}(attrs={})
-              get_proxy(associations[#{name.inspect}]).build(attrs)
+            def build_#{name}(attrs={}, &block)
+              get_proxy(associations[#{name.inspect}]).build(attrs, &block)
             end
 
-            def create_#{name}(attrs={})
-              get_proxy(associations[#{name.inspect}]).create(attrs)
+            def create_#{name}(attrs={}, &block)
+              get_proxy(associations[#{name.inspect}]).create(attrs, &block)
             end
 
-            def create_#{name}!(attrs={})
-              get_proxy(associations[#{name.inspect}]).create!(attrs)
+            def create_#{name}!(attrs={}, &block)
+              get_proxy(associations[#{name.inspect}]).create!(attrs, &block)
             end
           end_eval
         end

lib/mongo_mapper/plugins/keys.rb

@@ -169,6 +169,7 @@ def initialize(attrs={})
         @_new = true
         initialize_default_values(attrs)
         self.attributes = attrs
+        yield self if block_given?
       end
 
       def initialize_from_database(attrs={})

lib/mongo_mapper/plugins/querying.rb

@@ -24,11 +24,17 @@ def first_or_new(args)
         end
 
         def create(*docs)
-          initialize_each(*docs) { |doc| doc.save }
+          initialize_each(*docs) do |doc|
+            yield doc if block_given?
+            doc.save
+          end
         end
 
         def create!(*docs)
-          initialize_each(*docs) { |doc| doc.save! }
+          initialize_each(*docs) do |doc|
+            yield doc if block_given?
+            doc.save!
+          end
         end
 
         def update(*args)

spec/functional/associations/belongs_to_proxy_spec.rb

@@ -7,7 +7,6 @@
     @comment_class = Doc do
       key :post_id, ObjectId
     end
-
     @comment_class.belongs_to :post, :class => @post_class
   end
 
@@ -49,7 +48,7 @@ def post?
     comment.post.should == post
     comment.post.nil?.should be_false
   end
-  
+
   it "should not reload the association when replacing" do
     post = @post_class.new(:name => 'mongomapper')
     comment = @comment_class.new(:name => 'Foo!', :post => post)
@@ -59,10 +58,10 @@ def post?
   it "should properly assign the associated object when assigning the association with create" do
     child_class = Doc('Child')
     parent_class = Doc('Parent')
-    
+
     parent_class.one :child, :class => child_class
     child_class.belongs_to :parent, :class => parent_class
-    
+
     child = child_class.create(:parent => parent_class.create)
     child.parent.child.should == child
   end
@@ -160,49 +159,70 @@ class ::Thing
     end
   end
 
-  it "should be able to build" do
-    @comment_class.belongs_to :post, :class => @post_class
-
-    comment = @comment_class.create
-    post = comment.build_post(:title => 'Hello, world!')
-    comment.post.should be_instance_of(@post_class)
-    comment.post.should be_new
-    comment.post.title.should == 'Hello, world!'
-    comment.post.should == post
-    comment.post_id.should == post.id
-  end
-
-  it "should be able to create" do
-    @comment_class.belongs_to :post, :class => @post_class
+  context "when creating documents" do
+    let(:comment) { @comment_class.create }
 
-    comment = @comment_class.create
-    post = comment.create_post(:title => 'Hello, world!')
-    comment.post.should be_instance_of(@post_class)
-    comment.post.should_not be_new
-    comment.post.title.should == 'Hello, world!'
-    comment.post.should == post
-    comment.post_id.should == post.id
-  end
-
-  context "#create!" do
     before do
       @post_class.key :title, String, :required => true
       @comment_class.belongs_to :post, :class => @post_class
     end
 
-    it "should raise exception if invalid" do
-      comment = @comment_class.create
-      expect { comment.create_post! }.to raise_error(MongoMapper::DocumentNotValid)
+    context "#build" do
+      it "should work" do
+        post = comment.build_post(:title => 'Hello, world!')
+        comment.post.should be_instance_of(@post_class)
+        comment.post.should be_new
+        comment.post.title.should == 'Hello, world!'
+        comment.post.should == post
+        comment.post_id.should == post.id
+      end
+
+      it "should accept a block" do
+        comment.build_post(:title => 'Hello, world!') do |post|
+          post.title = "Hello world!"
+        end
+        comment.post.title.should == "Hello world!"
+      end
     end
 
-    it "should work if valid" do
-      comment = @comment_class.create
-      post = comment.create_post!(:title => 'Hello, world!')
-      comment.post.should be_instance_of(@post_class)
-      comment.post.should_not be_new
-      comment.post.title.should == 'Hello, world!'
-      comment.post.should == post
-      comment.post_id.should == post.id
+    context "#create" do
+      it "should work" do
+        post = comment.create_post(:title => 'Hello, world!')
+        comment.post.should be_instance_of(@post_class)
+        comment.post.should_not be_new
+        comment.post.title.should == 'Hello, world!'
+        comment.post.should == post
+        comment.post_id.should == post.id
+      end
+
+      it "should accept a block" do
+        comment.create_post(:title => 'Hello, world!') do |post|
+          post.title = "Hello world!"
+        end
+        comment.post.title.should == "Hello world!"
+      end
+    end
+
+    context "#create!" do
+      it "should accept a block" do
+        comment.create_post! do |post|
+          post.title = "Hello world!"
+        end
+        comment.post.title.should == "Hello world!"
+      end
+
+      it "should raise exception if invalid" do
+        expect { comment.create_post! }.to raise_error(MongoMapper::DocumentNotValid)
+      end
+
+      it "should work if valid" do
+        post = comment.create_post!(:title => 'Hello, world!')
+        comment.post.should be_instance_of(@post_class)
+        comment.post.should_not be_new
+        comment.post.title.should == 'Hello, world!'
+        comment.post.should == post
+        comment.post_id.should == post.id
+      end
     end
   end
 

spec/functional/associations/many_documents_proxy_spec.rb

@@ -290,7 +290,7 @@ def search_pets
     end
   end
 
-  context "build" do
+  context "#build" do
     it "should assign foreign key" do
       project = Project.create
       status = project.statuses.build
@@ -335,9 +335,17 @@ def search_pets
       status = project.statuses.build(:name => 'Foo')
       status.should be_new
     end
+
+    it "should accept a block" do
+      project = Project.new
+      status = project.statuses.build do |doc|
+        doc.name = "Foo"
+      end
+      project.statuses[0].name.should == "Foo"
+    end
   end
 
-  context "create" do
+  context "#create" do
     it "should assign foreign key" do
       project = Project.create
       status = project.statuses.create(:name => 'Foo!')
@@ -363,9 +371,17 @@ def search_pets
       project.statuses.create(:name => 'Foo!')
       project.statuses.size.should == 1
     end
+
+    it "should accept a block" do
+      project = Project.new
+      status = project.statuses.create do |doc|
+        doc.name = "Foo"
+      end
+      project.statuses.first.name.should == "Foo"
+    end
   end
 
-  context "create!" do
+  context "#create!" do
     it "should assign foreign key" do
       project = Project.create
       status = project.statuses.create!(:name => 'Foo!')
@@ -398,6 +414,14 @@ def search_pets
       project.statuses.create!(:name => 'Foo!')
       project.statuses.size.should == 1
     end
+
+    it "should accept a block" do
+      project = Project.new
+      status = project.statuses.create! do |doc|
+        doc.name = "Foo"
+      end
+      status.name.should == "Foo"
+    end
   end
 
   context "count" do