[pypy-commit] pypy py3.5: itertools.islice: use same logic as CPython and fix #2643

[rlamy]

Author: Ronan Lamy <ronan.lamy at gmail.com>
Branch: py3.5
Changeset: r92315:17c8c1d27c41
Date: 2017-09-04 03:32 +0100
http://bitbucket.org/pypy/pypy/changeset/17c8c1d27c41/

Log:	itertools.islice: use same logic as CPython and fix #2643

diff --git a/pypy/module/itertools/interp_itertools.py b/pypy/module/itertools/interp_itertools.py
--- a/pypy/module/itertools/interp_itertools.py
+++ b/pypy/module/itertools/interp_itertools.py
@@ -310,8 +310,7 @@
 
 islice_ignore_items_driver = jit.JitDriver(name='islice_ignore_items',
                                            greens=['tp'],
-                                           reds=['num', 'w_islice',
-                                                 'w_iterator'])
+                                           reds=['w_islice', 'w_iterator'])
 
 class W_ISlice(W_Root):
     def __init__(self, space, w_iterable, w_startstop, args_w):
@@ -325,7 +324,7 @@
             w_stop = w_startstop
         elif num_args <= 2:
             if space.is_w(w_startstop, space.w_None):
-                start = -1
+                start = 0
             else:
                 start = self.arg_int_w(w_startstop, 0,
                  "Indicies for islice() must be None or non-negative integers")
@@ -352,9 +351,10 @@
         else:
             step = 1
 
-        self.ignore = step - 1
-        self.start = start
+        self.count = 0
+        self.next = start
         self.stop = stop
+        self.step = step
 
     def arg_int_w(self, w_obj, minimum, errormsg):
         space = self.space
@@ -372,70 +372,52 @@
         return self
 
     def next_w(self):
-        if self.start >= 0:               # first call only
-            ignore = self.start
-            self.start = -1
-        else:                             # all following calls
-            ignore = self.ignore
-        stop = self.stop
-        if stop >= 0:
-            if stop <= ignore:
-                self.stop = 0   # reset the state so that a following next_w()
-                                # has no effect any more
-                if stop > 0:
-                    self._ignore_items(stop)
-                self.iterable = None
-                raise OperationError(self.space.w_StopIteration,
-                                     self.space.w_None)
-            self.stop = stop - (ignore + 1)
-        if ignore > 0:
-            self._ignore_items(ignore)
         if self.iterable is None:
             raise OperationError(self.space.w_StopIteration, self.space.w_None)
+        self._ignore_items()
+        stop = self.stop
+        if 0 <= stop <= self.count:
+            self.iterable = None
+            raise OperationError(self.space.w_StopIteration,
+                                    self.space.w_None)
         try:
-            return self.space.next(self.iterable)
+            item = self.space.next(self.iterable)
         except OperationError as e:
             if e.match(self.space, self.space.w_StopIteration):
                 self.iterable = None
             raise
+        self.count += 1
+        oldnext = self.next
+        self.next += self.step
+        if self.next < oldnext or self.next > stop >= 0:
+            self.next = stop
+        return item
 
-    def _ignore_items(self, num):
+    def _ignore_items(self):
         w_iterator = self.iterable
-        if w_iterator is None:
-            raise OperationError(self.space.w_StopIteration, self.space.w_None)
-
         tp = self.space.type(w_iterator)
         while True:
-            islice_ignore_items_driver.jit_merge_point(tp=tp,
-                                                       num=num,
-                                                       w_islice=self,
-                                                       w_iterator=w_iterator)
+            islice_ignore_items_driver.jit_merge_point(
+                tp=tp, w_islice=self, w_iterator=w_iterator)
+            if self.count >= self.next:
+                break
             try:
                 self.space.next(w_iterator)
             except OperationError as e:
                 if e.match(self.space, self.space.w_StopIteration):
                     self.iterable = None
                 raise
-            num -= 1
-            if num <= 0:
-                break
+            self.count += 1
 
     def descr_reduce(self, space):
         if self.iterable is None:
             return space.newtuple([
                 space.type(self),
                 space.newtuple([space.iter(space.newlist([])),
-                                space.newint(0),
-                                space.newint(0),
-                                space.newint(1),
-                            ]),
+                                space.newint(0)]),
+                space.newint(0),
             ])
-        start = self.start
         stop = self.stop
-        if start == -1:
-            w_start = space.w_None
-        else:
-            w_start = space.newint(start)
         if stop == -1:
             w_stop = space.w_None
         else:
@@ -443,11 +425,15 @@
         return space.newtuple([
             space.type(self),
             space.newtuple([self.iterable,
-                            w_start,
+                            space.newint(self.next),
                             w_stop,
-                            space.newint(self.ignore + 1)]),
+                            space.newint(self.step)]),
+            space.newint(self.count),
         ])
 
+    def descr_setstate(self, space, w_state):
+        self.count = space.int_w(w_state)
+
 def W_ISlice___new__(space, w_subtype, w_iterable, w_startstop, args_w):
     r = space.allocate_instance(W_ISlice, w_subtype)
     r.__init__(space, w_iterable, w_startstop, args_w)
diff --git a/pypy/module/itertools/test/test_itertools.py b/pypy/module/itertools/test/test_itertools.py
--- a/pypy/module/itertools/test/test_itertools.py
+++ b/pypy/module/itertools/test/test_itertools.py
@@ -214,6 +214,9 @@
         assert list(itertools.islice(range(10), None,None)) == list(range(10))
         assert list(itertools.islice(range(10), None,None,None)) == list(range(10))
 
+        it = itertools.islice([0, 1, 2], None, None, 2)
+        assert list(it) == [0, 2]
+
         import weakref
         for args in [(1,), (None,), (0, None, 2)]:
             it = (x for x in (1, 2, 3))