diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 8c13dc93..71e18932 100644 --- a/mongoengine/queryset.py +++ b/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: - raise StopIteration - return self._document._from_son(self._cursor.next()) + try: + if self._limit == 0: + 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. diff --git a/tests/queryset.py b/tests/queryset.py index b55c1876..e9d8e754 100644 --- a/tests/queryset.py +++ b/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. """