bpo-27212: Modify islice recipe to consume initial values preceding start (GH-6195)

csabella · rhettinger · commit da1734c58d2f · 2018-03-26T18:29:33.000-07:00
diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst
@@ -436,15 +436,24 @@ loops that truncate the stream.
           # islice('ABCDEFG', 2, None) --> C D E F G
           # islice('ABCDEFG', 0, None, 2) --> A C E G
           s = slice(*args)
-          it = iter(range(s.start or 0, s.stop or sys.maxsize, s.step or 1))
+          start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1
+          it = iter(range(start, stop, step))
           try:
               nexti = next(it)
           except StopIteration:
+              # Consume *iterable* up to the *start* position.
+              for i, element in zip(range(start), iterable):
+                  pass
               return
-          for i, element in enumerate(iterable):
-              if i == nexti:
-                  yield element
-                  nexti = next(it)
+          try:
+              for i, element in enumerate(iterable):
+                  if i == nexti:
+                      yield element
+                      nexti = next(it)
+          except StopIteration:
+              # Consume to *stop*.
+              for i, element in zip(range(i + 1, stop), iterable):
+                  pass
 
    If *start* is ``None``, then iteration starts at zero. If *step* is ``None``,
    then the step defaults to one.
@@ -688,8 +697,8 @@ which incur interpreter overhead.
        # tail(3, 'ABCDEFG') --> E F G
        return iter(collections.deque(iterable, maxlen=n))
 
-   def consume(iterator, n):
-       "Advance the iterator n-steps ahead. If n is none, consume entirely."
+   def consume(iterator, n=None):
+       "Advance the iterator n-steps ahead. If n is None, consume entirely."
        # Use functions that consume iterators at C speed.
        if n is None:
            # feed the entire iterator into a zero-length deque
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
@@ -1192,6 +1192,7 @@ def test_islice(self):
                 (10, 20, 3),
                 (10, 3, 20),
                 (10, 20),
+                (10, 10),
                 (10, 3),
                 (20,)
                 ]:
@@ -1218,6 +1219,10 @@ def test_islice(self):
         self.assertEqual(list(islice(it, 3)), list(range(3)))
         self.assertEqual(list(it), list(range(3, 10)))
 
+        it = iter(range(10))
+        self.assertEqual(list(islice(it, 3, 3)), [])
+        self.assertEqual(list(it), list(range(3, 10)))
+
         # Test invalid arguments
         ra = range(10)
         self.assertRaises(TypeError, islice, ra)
@@ -1604,6 +1609,48 @@ def test_takewhile(self):
         self.assertEqual(list(takewhile(lambda x: x<5, [1,4,6,4,1])), [1,4])
 
 
+class TestPurePythonRoughEquivalents(unittest.TestCase):
+
+    @staticmethod
+    def islice(iterable, *args):
+        s = slice(*args)
+        start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1
+        it = iter(range(start, stop, step))
+        try:
+            nexti = next(it)
+        except StopIteration:
+            # Consume *iterable* up to the *start* position.
+            for i, element in zip(range(start), iterable):
+                pass
+            return
+        try:
+            for i, element in enumerate(iterable):
+                if i == nexti:
+                    yield element
+                    nexti = next(it)
+        except StopIteration:
+            # Consume to *stop*.
+            for i, element in zip(range(i + 1, stop), iterable):
+                pass
+
+    def test_islice_recipe(self):
+        self.assertEqual(list(self.islice('ABCDEFG', 2)), list('AB'))
+        self.assertEqual(list(self.islice('ABCDEFG', 2, 4)), list('CD'))
+        self.assertEqual(list(self.islice('ABCDEFG', 2, None)), list('CDEFG'))
+        self.assertEqual(list(self.islice('ABCDEFG', 0, None, 2)), list('ACEG'))
+        # Test items consumed.
+        it = iter(range(10))
+        self.assertEqual(list(self.islice(it, 3)), list(range(3)))
+        self.assertEqual(list(it), list(range(3, 10)))
+        it = iter(range(10))
+        self.assertEqual(list(self.islice(it, 3, 3)), [])
+        self.assertEqual(list(it), list(range(3, 10)))
+        # Test that slice finishes in predictable state.
+        c = count()
+        self.assertEqual(list(self.islice(c, 1, 3, 50)), [1])
+        self.assertEqual(next(c), 3)
+
+
 class TestGC(unittest.TestCase):
 
     def makecycle(self, iterator, container):
@@ -2158,6 +2205,17 @@ def test_permutations_sizeof(self):
 ...     "Return function(0), function(1), ..."
 ...     return map(function, count(start))
 
+>>> import collections
+>>> def consume(iterator, n=None):
+...     "Advance the iterator n-steps ahead. If n is None, consume entirely."
+...     # Use functions that consume iterators at C speed.
+...     if n is None:
+...         # feed the entire iterator into a zero-length deque
+...         collections.deque(iterator, maxlen=0)
+...     else:
+...         # advance to the empty slice starting at position n
+...         next(islice(iterator, n, n), None)
+
 >>> def nth(iterable, n, default=None):
 ...     "Returns the nth item or a default value"
 ...     return next(islice(iterable, n, None), default)
@@ -2298,6 +2356,14 @@ def test_permutations_sizeof(self):
 >>> list(islice(tabulate(lambda x: 2*x), 4))
 [0, 2, 4, 6]
 
+>>> it = iter(range(10))
+>>> consume(it, 3)
+>>> next(it)
+3
+>>> consume(it)
+>>> next(it, 'Done')
+'Done'
+
 >>> nth('abcde', 3)
 'd'
 
@@ -2386,6 +2452,7 @@ def test_main(verbose=None):
     test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC,
                     RegressionTests, LengthTransparency,
                     SubclassWithKwargsTest, TestExamples,
+                    TestPurePythonRoughEquivalents,
                     SizeofTest)
     support.run_unittest(*test_classes)
 
diff --git a/Misc/NEWS.d/next/Documentation/2018-03-22-19-23-04.bpo-27212.wrE5KR.rst b/Misc/NEWS.d/next/Documentation/2018-03-22-19-23-04.bpo-27212.wrE5KR.rst
@@ -0,0 +1,2 @@
+Modify documentation for the :func:`islice` recipe to consume initial values
+up to the start index.

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	+Modify documentation for the :func:`islice` recipe to consume initial values
	`2`	`+up to the start index.`