Validate sort param for find_and_modify PYTHON-416

- find_and_modify now accepts a list of (key, direction) pairs
- passing SON or OrderedDict is still accepted but deprecated
- passing dict of len 1 is still accepted but deprecated
- passing dict of len > 1 raises TypeError

pymongo/collection.py

@@ -26,6 +26,13 @@ from pymongo.cursor import Cursor
 from pymongo.errors import ConfigurationError, InvalidName
 
 
+try:
+    from collections import OrderedDict
+    ordered_types = (SON, OrderedDict)
+except ImportError:
+    ordered_types = SON
+
+
 def _gen_index_name(keys):
     """Generate an index name from the set of fields it is over.
     """
@@ -1228,7 +1235,8 @@ class Collection(common.BaseObject):
         else:
             return res.get("results")
 
-    def find_and_modify(self, query={}, update=None, upsert=False, **kwargs):
+    def find_and_modify(self, query={}, update=None,
+                        upsert=False, sort=None, **kwargs):
         """Update and return an object.
 
         This is a thin wrapper around the findAndModify_ command. The
@@ -1243,13 +1251,15 @@ class Collection(common.BaseObject):
 
         :Parameters:
             - `query`: filter for the update (default ``{}``)
-            - `sort`: priority if multiple objects match (default ``{}``)
             - `update`: see second argument to :meth:`update` (no default)
+            - `upsert`: insert if object doesn't exist (default ``False``)
+            - `sort`: a list of (key, direction) pairs specifying the sort
+              order for this query. See :meth:`~pymongo.cursor.Cursor.sort`
+              for details.
             - `remove`: remove rather than updating (default ``False``)
             - `new`: return updated rather than original object
               (default ``False``)
             - `fields`: see second argument to :meth:`find` (default all)
-            - `upsert`: insert if object doesn't exist (default ``False``)
             - `**kwargs`: any other options the findAndModify_ command
               supports can be passed here.
 
@@ -1260,6 +1270,9 @@ class Collection(common.BaseObject):
 
         .. note:: Requires server version **>= 1.3.0**
 
+        .. versionchanged:: 2.3+
+           Deprecated the use of mapping types for the sort parameter
+
         .. versionadded:: 1.10
         """
         if (not update and not kwargs.get('remove', None)):
@@ -1275,6 +1288,22 @@ class Collection(common.BaseObject):
             kwargs['update'] = update
         if upsert:
             kwargs['upsert'] = upsert
+        if sort:
+            # Accept a list of tuples to match Cursor's sort parameter.
+            if isinstance(sort, list):
+                kwargs['sort'] = helpers._index_document(sort)
+            # Accept OrderedDict, SON, and dict with len == 1 so we
+            # don't break existing code already using find_and_modify.
+            elif (isinstance(sort, ordered_types) or
+                  isinstance(sort, dict) and len(sort) == 1):
+                warnings.warn("Passing mapping types for `sort` is deprecated,"
+                              " use a list of (key, direction) pairs instead",
+                              DeprecationWarning)
+                kwargs['sort'] = sort
+            else:
+                raise TypeError("sort must be a list of (key, direction) "
+                                 "pairs, a dict of len 1, or an instance of "
+                                 "SON or OrderedDict")
 
         no_obj_error = "No matching object found"
 

test/test_collection.py

@@ -1709,6 +1709,54 @@ class TestCollection(unittest.TestCase):
                                     as_class=ExtendedDict)
         self.assertTrue(isinstance(result, ExtendedDict))
 
+    def test_find_and_modify_with_sort(self):
+        c = self.db.test
+        c.drop()
+        for j in xrange(5):
+            c.insert({'j': j, 'i': 0}, safe=True)
+
+        sort={'j': DESCENDING}
+        self.assertEqual(4, c.find_and_modify({},
+                                              {'$inc': {'i': 1}},
+                                              sort=sort)['j'])
+        sort={'j': ASCENDING}
+        self.assertEqual(0, c.find_and_modify({},
+                                              {'$inc': {'i': 1}},
+                                              sort=sort)['j'])
+        sort=[('j', DESCENDING)]
+        self.assertEqual(4, c.find_and_modify({},
+                                              {'$inc': {'i': 1}},
+                                              sort=sort)['j'])
+        sort=[('j', ASCENDING)]
+        self.assertEqual(0, c.find_and_modify({},
+                                              {'$inc': {'i': 1}},
+                                              sort=sort)['j'])
+        sort=SON([('j', DESCENDING)])
+        self.assertEqual(4, c.find_and_modify({},
+                                              {'$inc': {'i': 1}},
+                                              sort=sort)['j'])
+        sort=SON([('j', ASCENDING)])
+        self.assertEqual(0, c.find_and_modify({},
+                                              {'$inc': {'i': 1}},
+                                              sort=sort)['j'])
+        try:
+            from collections import OrderedDict
+            sort=OrderedDict([('j', DESCENDING)])
+            self.assertEqual(4, c.find_and_modify({},
+                                                  {'$inc': {'i': 1}},
+                                                  sort=sort)['j'])
+            sort=OrderedDict([('j', ASCENDING)])
+            self.assertEqual(0, c.find_and_modify({},
+                                                  {'$inc': {'i': 1}},
+                                                  sort=sort)['j'])
+        except ImportError:
+            pass
+        # Test that a standard dict with two keys is rejected.
+        sort={'j': DESCENDING, 'foo': DESCENDING}
+        self.assertRaises(TypeError, c.find_and_modify, {},
+                                                         {'$inc': {'i': 1}},
+                                                         sort=sort)
+
     def test_find_with_nested(self):
         if not version.at_least(self.db.connection, (2, 0, 0)):
             raise SkipTest("nested $and and $or requires MongoDB >= 2.0")