Cleaner ordering code + allow clearing an explicit sort

mongoengine/queryset/base.py

@@ -971,13 +971,31 @@ class BaseQuerySet(object):
     def order_by(self, *keys):
         """Order the :class:`~mongoengine.queryset.QuerySet` by the keys. The
         order may be specified by prepending each of the keys by a + or a -.
-        Ascending order is assumed.
+        Ascending order is assumed. If no keys are passed, existing ordering
+        is cleared instead.
 
         :param keys: fields to order the query results by; keys may be
             prefixed with **+** or **-** to determine the ordering direction
         """
         queryset = self.clone()
-        queryset._ordering = queryset._get_order_by(keys)
+
+        old_ordering = queryset._ordering
+        new_ordering = queryset._get_order_by(keys)
+
+        if queryset._cursor_obj:
+
+            # If a cursor object has already been created, apply the sort to it
+            if new_ordering:
+                queryset._cursor_obj.sort(new_ordering)
+
+            # If we're trying to clear a previous explicit ordering, we need
+            # to clear the cursor entirely (because PyMongo doesn't allow
+            # clearing an existing sort on a cursor).
+            elif old_ordering:
+                queryset._cursor_obj = None
+
+        queryset._ordering = new_ordering
+
         return queryset
 
     def comment(self, text):
@@ -1423,10 +1441,13 @@ class BaseQuerySet(object):
             raise StopIteration
 
         raw_doc = self._cursor.next()
+
         if self._as_pymongo:
             return self._get_as_pymongo(raw_doc)
-        doc = self._document._from_son(raw_doc,
+
-                                       _auto_dereference=self._auto_dereference, only_fields=self.only_fields)
+        doc = self._document._from_son(
+            raw_doc, _auto_dereference=self._auto_dereference,
+            only_fields=self.only_fields)
 
         if self._scalar:
             return self._get_scalar(doc)
@@ -1436,7 +1457,6 @@ class BaseQuerySet(object):
     def rewind(self):
         """Rewind the cursor to its unevaluated state.
 
-
         .. versionadded:: 0.3
         """
         self._iter = False
@@ -1486,11 +1506,16 @@ class BaseQuerySet(object):
 
     @property
     def _cursor(self):
-        if self._cursor_obj is None:
+        """Return a PyMongo cursor object corresponding to this queryset."""
 
-            # In PyMongo 3+, we define the read preference on a collection
+        # If _cursor_obj already exists, return it immediately.
-            # level, not a cursor level. Thus, we need to get a cloned
+        if self._cursor_obj is not None:
-            # collection object using `with_options` first.
+            return self._cursor_obj
+
+        # Create a new PyMongo cursor.
+        # XXX In PyMongo 3+, we define the read preference on a collection
+        # level, not a cursor level. Thus, we need to get a cloned collection
+        # object using `with_options` first.
         if IS_PYMONGO_3 and self._read_preference is not None:
             self._cursor_obj = self._collection\
                 .with_options(read_preference=self._read_preference)\
@@ -1498,17 +1523,23 @@ class BaseQuerySet(object):
         else:
             self._cursor_obj = self._collection.find(self._query,
                                                      **self._cursor_args)
-            # Apply where clauses to cursor
+        # Apply "where" clauses to cursor
         if self._where_clause:
             where_clause = self._sub_js_fields(self._where_clause)
             self._cursor_obj.where(where_clause)
 
+        # Apply ordering to the cursor.
+        # XXX self._ordering can be equal to:
+        # * None if we didn't explicitly call order_by on this queryset.
+        # * A list of PyMongo-style sorting tuples.
+        # * An empty list if we explicitly called order_by() without any
+        #   arguments. This indicates that we want to clear the default
+        #   ordering.
         if self._ordering:
-                # Apply query ordering
+            # explicit ordering
             self._cursor_obj.sort(self._ordering)
         elif self._ordering is None and self._document._meta['ordering']:
-                # Otherwise, apply the ordering from the document model, unless
+            # default ordering
-                # it's been explicitly cleared via order_by with no arguments
             order = self._get_order_by(self._document._meta['ordering'])
             self._cursor_obj.sort(order)
 
@@ -1696,8 +1727,15 @@ class BaseQuerySet(object):
                     raise err
         return ret
 
+
     def _get_order_by(self, keys):
-        """Creates a list of order by fields"""
+        """Given a list of MongoEngine-style sort keys, return a list
+        of sorting tuples that can be applied to a PyMongo cursor. For
+        example:
+
+        >>> qs._get_order_by(['-last_name', 'first_name'])
+        [('last_name', -1), ('first_name', 1)]
+        """
         key_list = []
         for key in keys:
             if not key:
@@ -1710,17 +1748,19 @@ class BaseQuerySet(object):
             direction = pymongo.ASCENDING
             if key[0] == '-':
                 direction = pymongo.DESCENDING
+
             if key[0] in ('-', '+'):
                 key = key[1:]
+
             key = key.replace('__', '.')
             try:
                 key = self._document._translate_field_name(key)
             except Exception:
+                # TODO this exception should be more specific
                 pass
+
             key_list.append((key, direction))
 
-        if self._cursor_obj and key_list:
-            self._cursor_obj.sort(key_list)
         return key_list
 
     def _get_scalar(self, doc):

tests/queryset/queryset.py

@@ -1242,6 +1242,7 @@ class QuerySetTest(unittest.TestCase):
 
         BlogPost.drop_collection()
 
+        # default ordering should be used by default
         with db_ops_tracker() as q:
             BlogPost.objects.filter(title='whatever').first()
             self.assertEqual(len(q.get_ops()), 1)
@@ -1250,11 +1251,28 @@ class QuerySetTest(unittest.TestCase):
                 {'published_date': -1}
             )
 
+        # calling order_by() should clear the default ordering
         with db_ops_tracker() as q:
             BlogPost.objects.filter(title='whatever').order_by().first()
             self.assertEqual(len(q.get_ops()), 1)
             self.assertFalse('$orderby' in q.get_ops()[0]['query'])
 
+        # calling an explicit order_by should use a specified sort
+        with db_ops_tracker() as q:
+            BlogPost.objects.filter(title='whatever').order_by('published_date').first()
+            self.assertEqual(len(q.get_ops()), 1)
+            self.assertEqual(
+                q.get_ops()[0]['query']['$orderby'],
+                {'published_date': 1}
+            )
+
+        # calling order_by() after an explicit sort should clear it
+        with db_ops_tracker() as q:
+            qs = BlogPost.objects.filter(title='whatever').order_by('published_date')
+            qs.order_by().first()
+            self.assertEqual(len(q.get_ops()), 1)
+            self.assertFalse('$orderby' in q.get_ops()[0]['query'])
+
     def test_no_ordering_for_get(self):
         """ Ensure that Doc.objects.get doesn't use any ordering.
         """