Added full_result kwarg to update (#380)

docs/changelog.rst

@@ -2,6 +2,9 @@
 Changelog
 =========
 
+Changes in 0.8.3
+================
+- Added full_result kwarg to update (#380)
 
 Changes in 0.8.2
 ================

mongoengine/document.py

@@ -353,7 +353,13 @@ class Document(BaseDocument):
         been saved.
         """
         if not self.pk:
-            raise OperationError('attempt to update a document not yet saved')
+            if kwargs.get('upsert', False):
+                query = self.to_mongo()
+                if "_cls" in query:
+                    del(query["_cls"])
+                return self._qs.filter(**query).update_one(**kwargs)
+            else:
+                raise OperationError('attempt to update a document not yet saved')
 
         # Need to add shard key to query, or you get an error
         return self._qs.filter(**self._object_key).update_one(**kwargs)

mongoengine/queryset/queryset.py

@@ -474,7 +474,8 @@ class QuerySet(object):
 
         queryset._collection.remove(queryset._query, write_concern=write_concern)
 
-    def update(self, upsert=False, multi=True, write_concern=None, **update):
+    def update(self, upsert=False, multi=True, write_concern=None,
+               full_result=False, **update):
         """Perform an atomic update on the fields matched by the query.
 
         :param upsert: Any existing document with that "_id" is overwritten.
@@ -485,6 +486,8 @@ class QuerySet(object):
             ``save(..., write_concern={w: 2, fsync: True}, ...)`` will
             wait until at least two servers have recorded the write and
             will force an fsync on the primary server.
+        :param full_result: Return the full result rather than just the number
+            updated.
         :param update: Django-style update keyword arguments
 
         .. versionadded:: 0.2
@@ -506,12 +509,13 @@ class QuerySet(object):
                 update["$set"]["_cls"] = queryset._document._class_name
             else:
                 update["$set"] = {"_cls": queryset._document._class_name}
-
         try:
-            ret = queryset._collection.update(query, update, multi=multi,
-                                              upsert=upsert, **write_concern)
-            if ret is not None and 'n' in ret:
-                return ret['n']
+            result = queryset._collection.update(query, update, multi=multi,
+                                                 upsert=upsert, **write_concern)
+            if full_result:
+                return result
+            elif result:
+                return result['n']
         except pymongo.errors.OperationFailure, err:
             if unicode(err) == u'multi not coded yet':
                 message = u'update() method requires MongoDB 1.1.3+'

tests/queryset/queryset.py

@@ -536,6 +536,23 @@ class QuerySetTest(unittest.TestCase):
         self.assertEqual(club.members['John']['gender'], "F")
         self.assertEqual(club.members['John']['age'], 14)
 
+    def test_update_results(self):
+        self.Person.drop_collection()
+
+        result = self.Person(name="Bob", age=25).update(upsert=True, full_result=True)
+        self.assertIsInstance(result, dict)
+        self.assertTrue("upserted" in result)
+        self.assertFalse(result["updatedExisting"])
+
+        bob = self.Person.objects.first()
+        result = bob.update(set__age=30, full_result=True)
+        self.assertIsInstance(result, dict)
+        self.assertTrue(result["updatedExisting"])
+
+        self.Person(name="Bob", age=20).save()
+        result = self.Person.objects(name="Bob").update(set__name="bobby", multi=True)
+        self.assertEqual(result, 2)
+
     def test_upsert(self):
         self.Person.drop_collection()