Merge pull request #2335 from bagerard/fix_limit0_bug

Fix bug with Doc.objects.limit(0) which should return all docs

docs/changelog.rst

@@ -6,7 +6,8 @@ Changelog
 Development
 ===========
 - (Fill this out as you fix issues and develop your features).
-- Fixed a bug that made the queryset drop the read_preference after clone(). 
+- Fix a bug that made the queryset drop the read_preference after clone().
+- Fix the behavior of Doc.objects.limit(0) which should return all documents (similar to mongodb) #2311
 
 Changes in 0.20.0
 =================

mongoengine/queryset/base.py

@@ -83,6 +83,7 @@ class BaseQuerySet:
         self._cursor_obj = None
         self._limit = None
         self._skip = None
+
         self._hint = -1  # Using -1 as None is a valid value for hint
         self._collation = None
         self._batch_size = None
@@ -90,6 +91,13 @@ class BaseQuerySet:
         self._max_time_ms = None
         self._comment = None
 
+        # Hack - As people expect cursor[5:5] to return
+        # an empty result set. It's hard to do that right, though, because the
+        # server uses limit(0) to mean 'no limit'. So we set _empty
+        # in that case and check for it when iterating. We also unset
+        # it anytime we change _limit. Inspired by how it is done in pymongo.Cursor
+        self._empty = False
+
     def __call__(self, q_obj=None, **query):
         """Filter the selected documents by calling the
         :class:`~mongoengine.queryset.QuerySet` with a query.
@@ -162,6 +170,7 @@ class BaseQuerySet:
         [<User: User object>, <User: User object>]
         """
         queryset = self.clone()
+        queryset._empty = False
 
         # Handle a slice
         if isinstance(key, slice):
@@ -169,6 +178,8 @@ class BaseQuerySet:
             queryset._skip, queryset._limit = key.start, key.stop
             if key.start and key.stop:
                 queryset._limit = key.stop - key.start
+            if queryset._limit == 0:
+                queryset._empty = True
 
             # Allow further QuerySet modifications to be performed
             return queryset
@@ -394,7 +405,12 @@ class BaseQuerySet:
             :meth:`skip` that has been applied to this cursor into account when
             getting the count
         """
-        if self._limit == 0 and with_limit_and_skip is False or self._none:
+        if (
+            self._limit == 0
+            and with_limit_and_skip is False
+            or self._none
+            or self._empty
+        ):
             return 0
         count = self._cursor.count(with_limit_and_skip=with_limit_and_skip)
         self._cursor_obj = None
@@ -735,7 +751,9 @@ class BaseQuerySet:
         return doc_map
 
     def none(self):
-        """Helper that just returns a list"""
+        """Returns a queryset that never returns any objects and no query will be executed when accessing the results
+        inspired by django none() https://docs.djangoproject.com/en/dev/ref/models/querysets/#none
+        """
         queryset = self.clone()
         queryset._none = True
         return queryset
@@ -795,6 +813,7 @@ class BaseQuerySet:
             "_as_pymongo",
             "_limit",
             "_skip",
+            "_empty",
             "_hint",
             "_collation",
             "_auto_dereference",
@@ -835,6 +854,7 @@ class BaseQuerySet:
         """
         queryset = self.clone()
         queryset._limit = n
+        queryset._empty = False  # cancels the effect of empty
 
         # If a cursor object has already been created, apply the limit to it.
         if queryset._cursor_obj:
@@ -1586,7 +1606,7 @@ class BaseQuerySet:
     def __next__(self):
         """Wrap the result in a :class:`~mongoengine.Document` object.
         """
-        if self._limit == 0 or self._none:
+        if self._none or self._empty:
             raise StopIteration
 
         raw_doc = next(self._cursor)
@@ -1605,8 +1625,6 @@ class BaseQuerySet:
 
         return doc
 
-    next = __next__  # For Python2 support
-
     def rewind(self):
         """Rewind the cursor to its unevaluated state.
 

tests/queryset/test_queryset.py

@@ -114,6 +114,38 @@ class TestQueryset(unittest.TestCase):
         assert person.name == "User A"
         assert person.age == 20
 
+    def test_slicing_sets_empty_limit_skip(self):
+        self.Person.objects.insert(
+            [self.Person(name="User {}".format(i), age=i) for i in range(5)],
+            load_bulk=False,
+        )
+
+        self.Person.objects.create(name="User B", age=30)
+        self.Person.objects.create(name="User C", age=40)
+
+        qs = self.Person.objects()[1:2]
+        assert (qs._empty, qs._skip, qs._limit) == (False, 1, 1)
+        assert len(list(qs)) == 1
+
+        # Test edge case of [1:1] which should return nothing
+        # and require a hack so that it doesn't clash with limit(0)
+        qs = self.Person.objects()[1:1]
+        assert (qs._empty, qs._skip, qs._limit) == (True, 1, 0)
+        assert len(list(qs)) == 0
+
+        qs2 = qs[1:5]  # Make sure that further slicing resets _empty
+        assert (qs2._empty, qs2._skip, qs2._limit) == (False, 1, 4)
+        assert len(list(qs2)) == 4
+
+    def test_limit_0_returns_all_documents(self):
+        self.Person.objects.create(name="User A", age=20)
+        self.Person.objects.create(name="User B", age=30)
+
+        n_docs = self.Person.objects().count()
+
+        persons = list(self.Person.objects().limit(0))
+        assert len(persons) == 2 == n_docs
+
     def test_limit(self):
         """Ensure that QuerySet.limit works as expected."""
         user_a = self.Person.objects.create(name="User A", age=20)
@@ -377,6 +409,9 @@ class TestQueryset(unittest.TestCase):
 
         assert list(A.objects.none()) == []
         assert list(A.objects.none().all()) == []
+        assert list(A.objects.none().limit(1)) == []
+        assert list(A.objects.none().skip(1)) == []
+        assert list(A.objects.none()[:5]) == []
 
     def test_chaining(self):
         class A(Document):
@@ -4468,7 +4503,9 @@ class TestQueryset(unittest.TestCase):
         assert len(people) == 1
         assert people[0] == "User B"
 
-        people = list(self.Person.objects[1:1].scalar("name"))
+        # people = list(self.Person.objects[1:1].scalar("name"))
+        people = self.Person.objects[1:1]
+        people = people.scalar("name")
         assert len(people) == 0
 
         # Test slice out of range