Fix infinite loop possible when iterating columnbase types.

Fixes #2606
coleifer · Aug 18, 2022 · 3e3b195 · 3e3b195
1 parent eeff376
commit 3e3b195
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/peewee.py b/peewee.py
@@ -1240,6 +1240,7 @@ def __getitem__(self, item):
                                  'end-point.')
             return self.between(item.start, item.stop)
         return self == item
+    def __iter__(self): raise NotImplementedError  # Prevent accidents.
 
     def distinct(self):
         return NodeList((SQL('DISTINCT'), self))

diff --git a/tests/regressions.py b/tests/regressions.py
@@ -1736,3 +1736,8 @@ def assertNames(q, expected):
 
         q = DF.select().join(DFC).join(DFGC)
         assertNames(q.filter(dfc_set__dfgc_set__name='a1-1'), ['a'])
+
+
+class TestFunctionInfiniteLoop(BaseTestCase):
+    def test_function_infinite_loop(self):
+        self.assertRaises(NotImplementedError, lambda: list(fn.COUNT()))