lib/mongoo/cursor.rb in mongoo-0.3.1 vs lib/mongoo/cursor.rb in mongoo-0.4.0

@@ -9,31 +9,24 @@
       @mongo_cursor = mongo_cursor
     end
 
     def next_document
       if doc = @mongo_cursor.next_document
-        obj = @obj_class.new(doc, true)
-        Mongoo::IdentityMap.write(obj) if Mongoo::IdentityMap.on?
-        obj
+        obj_from_doc(doc)
       end
     end
 
     alias :next :next_document
 
     def each
       @mongo_cursor.each do |doc|
-        obj = @obj_class.new(doc, true)
-        Mongoo::IdentityMap.write(obj) if Mongoo::IdentityMap.on?
-        yield obj
+        yield obj_from_doc(doc)
       end
     end
 
     def to_a
-      arr = @mongo_cursor.to_a.collect { |doc| @obj_class.new(doc, true) }
-      if Mongoo::IdentityMap.on?
-        arr.each { |obj| Mongoo::IdentityMap.write(obj) }
-      end; arr
+      @mongo_cursor.to_a.collect { |doc| obj_from_doc(doc) }
     end
 
     def count
       @mongo_cursor.count
     end
@@ -62,9 +55,21 @@
       if @mongo_cursor.respond_to?(name)
         @mongo_cursor.send name, *args
       else
         super
       end
+    end
+  
+    def obj_from_doc(doc)
+      obj = nil
+      if Mongoo::IdentityMap.on?
+        if obj = Mongoo::IdentityMap.read(doc["_id"])
+          obj.merge!(doc)
+        end
+      end
+      obj ||= @obj_class.new(doc, true)
+      Mongoo::IdentityMap.write(obj) if Mongoo::IdentityMap.on?
+      obj
     end
 
   end
 end
\ No newline at end of file