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

This commit is contained in:
Harry Marr 2010-03-08 22:15:40 +00:00
parent 48e313fb44
commit ac3c857e1a
2 changed files with 24 additions and 3 deletions

View File

@ -416,9 +416,18 @@ class QuerySet(object):
def next(self):
"""Wrap the result in a :class:`~mongoengine.Document` object.
"""
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.

View File

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