Reduce groupby.__next__ calls in all_equal

Author: bbayles

more_itertools/recipes.py

@@ -222,6 +222,7 @@ def all_equal(iterable, key=None):
     for first in iterator:
         for second in iterator:
             return False
+        return True
     return True
 
 

tests/test_recipes.py

@@ -2,11 +2,12 @@
 from doctest import DocTestSuite
 from fractions import Fraction
 from functools import reduce
-from itertools import combinations, count, permutations
+from itertools import combinations, count, groupby, permutations
 from operator import mul
 from math import factorial
 from sys import version_info
 from unittest import TestCase, skipIf
+from unittest.mock import patch
 
 import more_itertools as mi
 
@@ -158,6 +159,22 @@ def test_key(self):
         self.assertTrue(mi.all_equal('4٤໔４৪', key=int))
         self.assertFalse(mi.all_equal('Abc', key=str.casefold))
 
+    @patch('more_itertools.recipes.groupby', autospec=True)
+    def test_groupby_calls(self, mock_groupby):
+        next_count = 0
+
+        class _groupby(groupby):
+            def __next__(true_self):
+                nonlocal next_count
+                next_count += 1
+                return super().__next__()
+
+        mock_groupby.side_effect = _groupby
+        iterable = iter('aaaaa')
+        self.assertTrue(mi.all_equal(iterable))
+        self.assertEqual(list(iterable), [])
+        self.assertEqual(next_count, 2)
+
 
 class QuantifyTests(TestCase):
     """Tests for ``quantify()``"""