From 219d316b4980dbec3192b82cd3cd9cf366b3f3ad Mon Sep 17 00:00:00 2001 From: Marc Tamlyn Date: Wed, 5 Oct 2011 13:26:57 +0100 Subject: [PATCH] Fix iteration on querysets. If iteration of a queryset was interrupted (by a break, or a caught error), the next iterator would start from the second element as the cursor had already moved to the first. This is fixed by adding a rewind into the __iter__ method. --- mongoengine/queryset.py | 1 + tests/queryset.py | 6 ++++++ 2 files changed, 7 insertions(+) diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index a6626855..a90d7591 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -1342,6 +1342,7 @@ class QuerySet(object): raise OperationError(u'Update failed [%s]' % unicode(e)) def __iter__(self): + self.rewind() return self def _sub_js_fields(self, code): diff --git a/tests/queryset.py b/tests/queryset.py index f093f6ab..5c7afe68 100644 --- a/tests/queryset.py +++ b/tests/queryset.py @@ -567,7 +567,13 @@ class QuerySetTest(unittest.TestCase): people1 = [person for person in queryset] people2 = [person for person in queryset] + # Check that it still works even if iteration is interrupted. + for person in queryset: + break + people3 = [person for person in queryset] + self.assertEqual(people1, people2) + self.assertEqual(people1, people3) def test_repr_iteration(self): """Ensure that QuerySet __repr__ can handle loops