reimplement itertools.groupby based on PyPy and add test

Author: timfel

graalpython/com.oracle.graal.python.test/src/tests/test_itertools.py

@@ -6,6 +6,7 @@
 import unittest
 from itertools import *
 
+
 class CombinationsTests(unittest.TestCase):
 
     def test_combinations_with_replacement(self):
@@ -24,3 +25,76 @@ def test_combinations_with_replacement(self):
             compare += a;
         self.assertEqual(result,compare)
 
+
+    def test_groupby(self):
+        # Check whether it accepts arguments correctly
+        self.assertEqual([], list(groupby([])))
+        self.assertEqual([], list(groupby([], key=id)))
+        self.assertRaises(TypeError, list, groupby('abc', []))
+        self.assertRaises(TypeError, groupby, None)
+        self.assertRaises(TypeError, groupby, 'abc', lambda x:x, 10)
+
+        # Check normal input
+        s = [(0, 10, 20), (0, 11,21), (0,12,21), (1,13,21), (1,14,22),
+             (2,15,22), (3,16,23), (3,17,23)]
+        dup = []
+        for k, g in groupby(s, lambda r:r[0]):
+            for elem in g:
+                self.assertEqual(k, elem[0])
+                dup.append(elem)
+        self.assertEqual(s, dup)
+
+        # Exercise pipes and filters style
+        s = 'abracadabra'
+        # sort s | uniq
+        r = [k for k, g in groupby(sorted(s))]
+        self.assertEqual(r, ['a', 'b', 'c', 'd', 'r'])
+        # sort s | uniq -d
+        r = [k for k, g in groupby(sorted(s)) if list(islice(g,1,2))]
+        self.assertEqual(r, ['a', 'b', 'r'])
+        # sort s | uniq -c
+        r = [(len(list(g)), k) for k, g in groupby(sorted(s))]
+        self.assertEqual(r, [(5, 'a'), (2, 'b'), (1, 'c'), (1, 'd'), (2, 'r')])
+        # sort s | uniq -c | sort -rn | head -3
+        r = sorted([(len(list(g)) , k) for k, g in groupby(sorted(s))], reverse=True)[:3]
+        self.assertEqual(r, [(5, 'a'), (2, 'r'), (2, 'b')])
+
+        # iter.__next__ failure
+        class ExpectedError(Exception):
+            pass
+        def delayed_raise(n=0):
+            for i in range(n):
+                yield 'yo'
+            raise ExpectedError
+        def gulp(iterable, keyp=None, func=list):
+            return [func(g) for k, g in groupby(iterable, keyp)]
+
+        # iter.__next__ failure on outer object
+        self.assertRaises(ExpectedError, gulp, delayed_raise(0))
+        # iter.__next__ failure on inner object
+        self.assertRaises(ExpectedError, gulp, delayed_raise(1))
+
+        # __eq__ failure
+        class DummyCmp:
+            def __eq__(self, dst):
+                raise ExpectedError
+        s = [DummyCmp(), DummyCmp(), None]
+
+        # __eq__ failure on outer object
+        self.assertRaises(ExpectedError, gulp, s, func=id)
+        # __eq__ failure on inner object
+        self.assertRaises(ExpectedError, gulp, s)
+
+        # keyfunc failure
+        def keyfunc(obj):
+            if keyfunc.skip > 0:
+                keyfunc.skip -= 1
+                return obj
+            else:
+                raise ExpectedError
+
+        # keyfunc failure on outer object
+        keyfunc.skip = 0
+        self.assertRaises(ExpectedError, gulp, [None], keyfunc)
+        keyfunc.skip = 1
+        self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)

graalpython/lib-graalpython/itertools.py

@@ -489,27 +489,66 @@ class groupby(object):
            uniquekeys.append(k)
     """
     def __init__(self, iterable, key=None):
-        if key is None:
-            key = lambda x: x
+        self._iterator = iter(iterable)
         self._keyfunc = key
-        self._iter = iter(iterable)
-        self._tgtkey = self._currkey = self._currvalue = range(0)
+        self._tgtkey = self._currkey = self._currvalue = None
 
     def __iter__(self):
         return self
 
     def __next__(self):
-        while self._currkey == self._tgtkey:
-            self._currvalue = next(self._iter) # Exit on StopIteration
-            self._currkey = self._keyfunc(self._currvalue)
-        self._tgtkey = self._currkey
-        return (self._currkey, self._grouper(self._tgtkey))
-
-    def _grouper(self, tgtkey):
-        while self._currkey == tgtkey:
-            yield self._currvalue
-            self._currvalue = next(self._iter) # Exit on StopIteration
-            self._currkey = self._keyfunc(self._currvalue)
+        self._skip_to_next_iteration_group()
+        key = self._tgtkey = self._currkey
+        grouper = _groupby(self, key)
+        return (key, grouper)
+
+    def _skip_to_next_iteration_group(self):
+        while True:
+            if self._currkey is None:
+                pass
+            elif self._tgtkey is None:
+                break
+            else:
+                if not self._tgtkey == self._currkey:
+                    break
+
+            newvalue = next(self._iterator)
+            if self._keyfunc is None:
+                newkey = newvalue
+            else:
+                newkey = self._keyfunc(newvalue)
+
+            self._currkey = newkey
+            self._currvalue = newvalue
+
+
+class _groupby():
+    def __init__(self, groupby, tgtkey):
+        self.groupby = groupby
+        self.tgtkey = tgtkey
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        groupby = self.groupby
+        if groupby._currvalue is None:
+            newvalue = next(groupby._iterator)
+            if groupby._keyfunc is None:
+                newkey = newvalue
+            else:
+                newkey = groupby._keyfunc(newvalue)
+            assert groupby._currvalue is None
+            groupby._currkey = newkey
+            groupby._currvalue = newvalue
+
+        assert groupby._currkey is not None
+        if not self.tgtkey == groupby._currkey:
+            raise StopIteration(None)
+        result = groupby._currvalue
+        groupby._currvalue = None
+        groupby._currkey = None
+        return result
 
 
 class combinations():