Added rewind to QuerySet, which is implicitly called when iteration finishes

mongoengine/queryset.py

@@ -416,9 +416,18 @@ class QuerySet(object):
     def next(self):
         """Wrap the result in a :class:`~mongoengine.Document` object.
         """
-        if self._limit == 0:
+        try:
-            raise StopIteration
+            if self._limit == 0:
-        return self._document._from_son(self._cursor.next())
+                raise StopIteration
+            return self._document._from_son(self._cursor.next())
+        except StopIteration, e:
+            self.rewind()
+            raise e
+
+    def rewind(self):
+        """Rewind the cursor to its unevaluated state.
+        """
+        self._cursor.rewind()
 
     def count(self):
         """Count the selected elements in the query.

tests/queryset.py

@@ -186,6 +186,18 @@ class QuerySetTest(unittest.TestCase):
         person = self.Person.objects.get(age=50)
         self.assertEqual(person.name, "User C")
 
+    def test_repeated_iteration(self):
+        """Ensure that QuerySet rewinds itself one iteration finishes.
+        """
+        self.Person(name='Person 1').save()
+        self.Person(name='Person 2').save()
+
+        queryset = self.Person.objects
+        people1 = [person for person in queryset]
+        people2 = [person for person in queryset]
+
+        self.assertEqual(people1, people2)
+
     def test_regex_query_shortcuts(self):
         """Ensure that contains, startswith, endswith, etc work.
         """