From 5ccd02653a986de07f0830259f27c80e9b8de387 Mon Sep 17 00:00:00 2001
From: Sergey Azovskov <last-g@yandex-team.ru>
Date: Thu, 5 Jun 2014 18:42:27 +0600
Subject: [PATCH] Added support for manipulate param in find_and_modify for
 consistency with find method

---
 pymongo/collection.py   | 11 +++++++++--
 test/test_collection.py | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/pymongo/collection.py b/pymongo/collection.py
index a2368db99..46ded9de8 100644
--- a/pymongo/collection.py
+++ b/pymongo/collection.py
@@ -1622,7 +1622,8 @@ class Collection(common.BaseObject):
             return res.get("results")
 
     def find_and_modify(self, query={}, update=None,
-                        upsert=False, sort=None, full_response=False, **kwargs):
+                        upsert=False, sort=None, full_response=False,
+                        manipulate=True, **kwargs):
         """Update and return an object.
 
         This is a thin wrapper around the findAndModify_ command. The
@@ -1654,6 +1655,9 @@ class Collection(common.BaseObject):
             - `new`: return updated rather than original object
               (default ``False``)
             - `fields`: see second argument to :meth:`find` (default all)
+            - `manipulate`: (optional): If True (the default), apply any
+               outgoing SON manipulators before returning. Do not works when
+               `full_response` is set to True.
             - `**kwargs`: any other options the findAndModify_ command
               supports can be passed here.
 
@@ -1720,7 +1724,10 @@ class Collection(common.BaseObject):
         if full_response:
             return out
         else:
-            return out.get('value')
+            document = out.get('value')
+            if manipulate:
+                document = self.__database._fix_outgoing(document, self)
+            return document
 
     def __iter__(self):
         return self
diff --git a/test/test_collection.py b/test/test_collection.py
index e557e0bd3..248d0ae2c 100644
--- a/test/test_collection.py
+++ b/test/test_collection.py
@@ -2350,6 +2350,43 @@ class TestCollection(unittest.TestCase):
         for doc in c.find(compile_re=False):
             self.assertTrue(isinstance(doc['r'], Regex))
 
+    def test_find_and_modify_with_manipulator(self):
+        class AddCollectionNameManipulator(SONManipulator):
+            def will_copy(self):
+                return True
+
+            def transform_incoming(self, son, collection):
+                copy = SON(son)
+                if 'collection' in copy:
+                    del copy['collection']
+                return copy
+
+            def transform_outgoing(self, son, collection):
+                copy = SON(son)
+                copy['collection'] = collection.name
+                return copy
+
+        self.db.add_son_manipulator(AddCollectionNameManipulator())
+
+        c = self.db.test
+        c.drop()
+        c.insert({'_id': 1, 'i': 1})
+
+        # Test correct findAndModify
+        # With manipulators
+        self.assertEqual({'_id': 1, 'i': 1, 'collection': 'test'},
+                         c.find_and_modify({'_id': 1}, {'$inc': {'i': 1}}))
+        self.assertEqual({'_id': 1, 'i': 3, 'collection': 'test'},
+                         c.find_and_modify({'_id': 1}, {'$inc': {'i': 1}},
+                                           new=True))
+        # With out manipulators
+        self.assertEqual({'_id': 1, 'i': 3},
+                         c.find_and_modify({'_id': 1}, {'$inc': {'i': 1}},
+                                           manipulate=False))
+        self.assertEqual({'_id': 1, 'i': 5},
+                         c.find_and_modify({'_id': 1}, {'$inc': {'i': 1}},
+                                           new=True, manipulate=False))
+
 
 if __name__ == "__main__":
     unittest.main()