[pypy-svn] r56493 - in pypy/dist/pypy/module/itertools: . test

[jlg at codespeak.net]

Author: jlg
Date: Sat Jul 12 15:20:03 2008
New Revision: 56493

Modified:
   pypy/dist/pypy/module/itertools/interp_itertools.py
   pypy/dist/pypy/module/itertools/test/test_itertools.py
Log:
interp_itertools - groupby consumes previous group when when next() called

Modified: pypy/dist/pypy/module/itertools/interp_itertools.py
==============================================================================
--- pypy/dist/pypy/module/itertools/interp_itertools.py	(original)
+++ pypy/dist/pypy/module/itertools/interp_itertools.py	Sat Jul 12 15:20:03 2008
@@ -731,7 +731,8 @@
         self.lookahead = False
         self.exhausted = False
         self.started = False
-        self.group_edge = True
+        # new_group - new group not started yet, next should not skip any items
+        self.new_group = True 
         self.w_lookahead = self.space.w_None
         self.w_key = self.space.w_None
 
@@ -742,6 +743,17 @@
         if self.exhausted:
             raise OperationError(self.space.w_StopIteration, self.space.w_None)
 
+        if not self.new_group:
+            # Consume unwanted input until we reach the next group
+            try:
+                while True:
+                    self.group_next(self.index)
+
+            except StopIteration:
+                pass
+            if self.exhausted:
+                raise OperationError(self.space.w_StopIteration, self.space.w_None)
+
         if not self.started:
             self.started = True
             try:
@@ -758,20 +770,11 @@
                     self.w_key = self.space.call_function(self.w_fun, w_obj)
                 self.lookahead = True
 
-        if not self.group_edge:
-            # Consume unwanted input until we reach the next group
-            try:
-                while True:
-                    self.group_next(self.index)
-            except StopIteration:
-                pass
-            if self.exhausted:
-                raise OperationError(self.space.w_StopIteration, self.space.w_None)
+        self.new_group = False
         w_iterator = self.space.wrap(W_GroupByIterator(self.space, self.index, self))
         return self.space.newtuple([self.w_key, w_iterator])
 
     def group_next(self, group_index):
-        self.group_edge = False
         if group_index < self.index:
             raise StopIteration
         else:
@@ -799,7 +802,7 @@
                     self.w_lookahead = w_obj
                     self.w_key = w_new_key
                     self.lookahead = True
-                    self.group_edge = True
+                    self.new_group = True #new group
                     raise StopIteration
 
 def W_GroupBy___new__(space, w_subtype, w_iterable, w_fun=None):

Modified: pypy/dist/pypy/module/itertools/test/test_itertools.py
==============================================================================
--- pypy/dist/pypy/module/itertools/test/test_itertools.py	(original)
+++ pypy/dist/pypy/module/itertools/test/test_itertools.py	Sat Jul 12 15:20:03 2008
@@ -434,13 +434,31 @@
             raises(StopIteration, g.next)
         raises(StopIteration, it.next)
 
+        # consumes after group started
         it = itertools.groupby([0, 0, 0, 0, 1])
         k1, g1 = it.next()
+        assert g1.next() == 0
         k2, g2 = it.next()
         raises(StopIteration, g1.next)
         assert g2.next() == 1
         raises(StopIteration, g2.next)
 
+        # skips with not started group
+        it = itertools.groupby([0, 0, 1])
+        k1, g1 = it.next()
+        k2, g2 = it.next()
+        raises(StopIteration, g1.next)
+        assert g2.next() == 1
+        raises(StopIteration, g2.next)
+
+        it = itertools.groupby([0, 1, 2])
+        k1, g1 = it.next()
+        k2, g2 = it.next()
+        k2, g3 = it.next()
+        raises(StopIteration, g1.next)
+        raises(StopIteration, g2.next)
+        assert g3.next() == 2
+
         def half_floor(x):
             return x // 2
         it = itertools.groupby([0, 1, 2, 3, 4, 5], half_floor)