QuerySet.modify() method to provide find_and_modify() like behaviour

docs/guide/querying.rst

@@ -488,8 +488,9 @@ calling it with keyword arguments::
 Atomic updates
 ==============
 Documents may be updated atomically by using the
-:meth:`~mongoengine.queryset.QuerySet.update_one` and
-:meth:`~mongoengine.queryset.QuerySet.update` methods on a
+:meth:`~mongoengine.queryset.QuerySet.update_one`,
+:meth:`~mongoengine.queryset.QuerySet.update` and
+:meth:`~mongoengine.queryset.QuerySet.modify` methods on a
 :meth:`~mongoengine.queryset.QuerySet`. There are several different "modifiers"
 that you may use with these methods:
 

mongoengine/queryset/base.py

@@ -10,6 +10,7 @@
 from bson.code import Code
 from bson import json_util
 import pymongo
+import pymongo.errors
 from pymongo.common import validate_read_preference
 
 from mongoengine import signals
@@ -484,6 +485,59 @@ def update_one(self, upsert=False, write_concern=None, **update):
         return self.update(
             upsert=upsert, multi=False, write_concern=write_concern, **update)
 
+    def modify(self, upsert=False, full_response=False, remove=False, new=False, **update):
+        """Update and return the updated document.
+
+        Returns either the document before or after modification based on `new`
+        parameter. If no documents match the query and `upsert` is false,
+        returns ``None``. If upserting and `new` is false, returns ``None``.
+
+        If the full_response parameter is ``True``, the return value will be
+        the entire response object from the server, including the 'ok' and
+        'lastErrorObject' fields, rather than just the modified document.
+        This is useful mainly because the 'lastErrorObject' document holds
+        information about the command's execution.
+
+        :param upsert: insert if document doesn't exist (default ``False``)
+        :param full_response: return the entire response object from the
+            server (default ``False``)
+        :param remove: remove rather than updating (default ``False``)
+        :param new: return updated rather than original document
+            (default ``False``)
+        :param update: Django-style update keyword arguments
+
+        .. versionadded:: 0.9
+        """
+
+        if remove and new:
+            raise OperationError("Conflicting parameters: remove and new")
+
+        if not update and not upsert and not remove:
+            raise OperationError("No update parameters, must either update or remove")
+
+        queryset = self.clone()
+        query = queryset._query
+        update = transform.update(queryset._document, **update)
+        sort = queryset._ordering
+
+        try:
+            result = queryset._collection.find_and_modify(
+                query, update, upsert=upsert, sort=sort, remove=remove, new=new,
+                full_response=full_response, **self._cursor_args)
+        except pymongo.errors.DuplicateKeyError, err:
+            raise NotUniqueError(u"Update failed (%s)" % err)
+        except pymongo.errors.OperationFailure, err:
+            raise OperationError(u"Update failed (%s)" % err)
+
+        if full_response:
+            if result["value"] is not None:
+                result["value"] = self._document._from_son(result["value"])
+        else:
+            if result is not None:
+                result = self._document._from_son(result)
+
+        return result
+
     def with_id(self, object_id):
         """Retrieve the object matching the id provided.  Uses `object_id` only
         and raises InvalidQueryError if a filter has been applied. Returns

tests/queryset/__init__.py

@@ -3,3 +3,4 @@
 from queryset import *
 from visitor import *
 from geo import *
+from modify import *

tests/queryset/modify.py

@@ -0,0 +1,102 @@
+import sys
+sys.path[0:0] = [""]
+
+import unittest
+
+from mongoengine import connect, Document, IntField
+
+__all__ = ("FindAndModifyTest",)
+
+
+class Doc(Document):
+    id = IntField(primary_key=True)
+    value = IntField()
+
+
+class FindAndModifyTest(unittest.TestCase):
+
+    def setUp(self):
+        connect(db="mongoenginetest")
+        Doc.drop_collection()
+
+    def assertDbEqual(self, docs):
+        self.assertEqual(list(Doc._collection.find().sort("id")), docs)
+
+    def test_modify(self):
+        Doc(id=0, value=0).save()
+        doc = Doc(id=1, value=1).save()
+
+        old_doc = Doc.objects(id=1).modify(set__value=-1)
+        self.assertEqual(old_doc.to_json(), doc.to_json())
+        self.assertDbEqual([{"_id": 0, "value": 0}, {"_id": 1, "value": -1}])
+
+    def test_modify_with_new(self):
+        Doc(id=0, value=0).save()
+        doc = Doc(id=1, value=1).save()
+
+        new_doc = Doc.objects(id=1).modify(set__value=-1, new=True)
+        doc.value = -1
+        self.assertEqual(new_doc.to_json(), doc.to_json())
+        self.assertDbEqual([{"_id": 0, "value": 0}, {"_id": 1, "value": -1}])
+
+    def test_modify_not_existing(self):
+        Doc(id=0, value=0).save()
+        self.assertEqual(Doc.objects(id=1).modify(set__value=-1), None)
+        self.assertDbEqual([{"_id": 0, "value": 0}])
+
+    def test_modify_with_upsert(self):
+        Doc(id=0, value=0).save()
+        old_doc = Doc.objects(id=1).modify(set__value=1, upsert=True)
+        self.assertEqual(old_doc, None)
+        self.assertDbEqual([{"_id": 0, "value": 0}, {"_id": 1, "value": 1}])
+
+    def test_modify_with_upsert_existing(self):
+        Doc(id=0, value=0).save()
+        doc = Doc(id=1, value=1).save()
+
+        old_doc = Doc.objects(id=1).modify(set__value=-1, upsert=True)
+        self.assertEqual(old_doc.to_json(), doc.to_json())
+        self.assertDbEqual([{"_id": 0, "value": 0}, {"_id": 1, "value": -1}])
+
+    def test_modify_with_upsert_with_new(self):
+        Doc(id=0, value=0).save()
+        new_doc = Doc.objects(id=1).modify(upsert=True, new=True, set__value=1)
+        self.assertEqual(new_doc.to_mongo(), {"_id": 1, "value": 1})
+        self.assertDbEqual([{"_id": 0, "value": 0}, {"_id": 1, "value": 1}])
+
+    def test_modify_with_remove(self):
+        Doc(id=0, value=0).save()
+        doc = Doc(id=1, value=1).save()
+
+        old_doc = Doc.objects(id=1).modify(remove=True)
+        self.assertEqual(old_doc.to_json(), doc.to_json())
+        self.assertDbEqual([{"_id": 0, "value": 0}])
+
+    def test_find_and_modify_with_remove_not_existing(self):
+        Doc(id=0, value=0).save()
+        self.assertEqual(Doc.objects(id=1).modify(remove=True), None)
+        self.assertDbEqual([{"_id": 0, "value": 0}])
+
+    def test_modify_with_order_by(self):
+        Doc(id=0, value=3).save()
+        Doc(id=1, value=2).save()
+        Doc(id=2, value=1).save()
+        doc = Doc(id=3, value=0).save()
+
+        old_doc = Doc.objects().order_by("-id").modify(set__value=-1)
+        self.assertEqual(old_doc.to_json(), doc.to_json())
+        self.assertDbEqual([
+            {"_id": 0, "value": 3}, {"_id": 1, "value": 2},
+            {"_id": 2, "value": 1}, {"_id": 3, "value": -1}])
+
+    def test_modify_with_fields(self):
+        Doc(id=0, value=0).save()
+        Doc(id=1, value=1).save()
+
+        old_doc = Doc.objects(id=1).only("id").modify(set__value=-1)
+        self.assertEqual(old_doc.to_mongo(), {"_id": 1})
+        self.assertDbEqual([{"_id": 0, "value": 0}, {"_id": 1, "value": -1}])
+
+
+if __name__ == '__main__':
+    unittest.main()