diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 3678b099..034b8476 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -98,18 +98,32 @@ class QuerySet(object): return self._cursor.count() def limit(self, n): - """Limit the number of returned documents to `n`. + """Limit the number of returned documents to `n`. This may also be + achieved using array-slicing syntax (e.g. ``User.objects[:5]``). """ self._cursor.limit(n) # Return self to allow chaining return self def skip(self, n): - """Skip `n` documents before returning the results. + """Skip `n` documents before returning the results. This may also be + achieved using array-slicing syntax (e.g. ``User.objects[5:]``). """ self._cursor.skip(n) return self + def __getitem__(self, key): + """Support skip and limit using getitem and slicing syntax. + """ + # Slice provided + if isinstance(key, slice): + self._cursor_obj = self._cursor[key] + # Allow further QuerySet modifications to be performed + return self + # Integer index provided + elif isinstance(key, int): + return self._document._from_son(self._cursor[key]) + 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 -. diff --git a/tests/queryset.py b/tests/queryset.py index d40e73b0..b73bab17 100644 --- a/tests/queryset.py +++ b/tests/queryset.py @@ -77,6 +77,26 @@ class QuerySetTest(unittest.TestCase): self.assertEqual(len(people), 1) self.assertEqual(people[0].name, 'User B') + person3 = self.Person(name="User C", age=40) + person3.save() + + # Test slice limit + people = list(self.Person.objects[:2]) + self.assertEqual(len(people), 2) + self.assertEqual(people[0].name, 'User A') + self.assertEqual(people[1].name, 'User B') + + # Test slice skip + people = list(self.Person.objects[1:]) + self.assertEqual(len(people), 2) + self.assertEqual(people[0].name, 'User B') + self.assertEqual(people[1].name, 'User C') + + # Test slice limit and skip + people = list(self.Person.objects[1:2]) + self.assertEqual(len(people), 1) + self.assertEqual(people[0].name, 'User B') + def test_find_one(self): """Ensure that a query using find_one returns a valid result. """ @@ -97,6 +117,15 @@ class QuerySetTest(unittest.TestCase): person = self.Person.objects(age__lt=30).first() self.assertEqual(person.name, "User A") + + # Use array syntax + person = self.Person.objects[0] + self.assertEqual(person.name, "User A") + + person = self.Person.objects[1] + self.assertEqual(person.name, "User B") + + self.assertRaises(IndexError, self.Person.objects.__getitem__, 2) # Find a document using just the object id person = self.Person.objects.with_id(person1.id)