Allow iterating over result wrapper for multiple times (#84)

05bit · Jan 16, 2018 · fcbb866 · fcbb866
1 parent b94a662
commit fcbb866
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/peewee_async.py b/peewee_async.py
@@ -796,8 +796,7 @@ def __init__(self, *, cursor=None, query=None):
         self._result_wrapper = self._get_result_wrapper(query)
 
     def __iter__(self):
-        while True:
-            yield self._result_wrapper.iterate()
+        return iter(self._result_wrapper)
 
     def __len__(self):
         return len(self._rows)

diff --git a/tests/__init__.py b/tests/__init__.py
@@ -593,6 +593,18 @@ def test(objects):
 
         self.run_with_managers(test)
 
+    def test_multiple_iterate_over_result(self):
+        @asyncio.coroutine
+        def test(objects):
+            obj1 = yield from objects.create(TestModel, text="Test 1")
+            obj2 = yield from objects.create(TestModel, text="Test 2")
+            result = yield from objects.execute(
+                TestModel.select().order_by(TestModel.text))
+            self.assertEqual(list(result), [obj1, obj2])
+            self.assertEqual(list(result), [obj1, obj2])
+
+        self.run_with_managers(test)
+
     def test_insert_many_rows_query(self):
         @asyncio.coroutine
         def test(objects):