Fix a bug in limit0 #2311

This commit is contained in:
Bastien Gérard 2020-05-26 23:37:55 +02:00
parent dc7da5204f
commit adb5f74ddb
2 changed files with 51 additions and 6 deletions

View File

@ -83,6 +83,7 @@ class BaseQuerySet:
self._cursor_obj = None
self._limit = None
self._skip = None
self._empty = False
self._hint = -1 # Using -1 as None is a valid value for hint
self._collation = None
self._batch_size = None
@ -162,6 +163,7 @@ class BaseQuerySet:
[<User: User object>, <User: User object>]
"""
queryset = self.clone()
queryset._empty = False
# Handle a slice
if isinstance(key, slice):
@ -169,6 +171,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 +398,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 +744,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
@ -794,6 +805,7 @@ class BaseQuerySet:
"_as_pymongo",
"_limit",
"_skip",
"_empty",
"_hint",
"_collation",
"_auto_dereference",
@ -834,6 +846,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:
@ -1584,7 +1597,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)
@ -1603,8 +1616,6 @@ class BaseQuerySet:
return doc
next = __next__ # For Python2 support
def rewind(self):
"""Rewind the cursor to its unevaluated state.

View File

@ -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)
@ -4442,7 +4474,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