Added array-style index/slicing syntax to QuerySet

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 -.

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.
         """
@@ -98,6 +118,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)
         self.assertEqual(person.name, "User A")