[Python-checkins] bpo-44571: Add itertool recipe for a variant of takewhile() (GH-28167)

[miss-islington]

https://github.com/python/cpython/commit/656b0bdfaae3a36d386afe3f7b991744528c3ff7
commit: 656b0bdfaae3a36d386afe3f7b991744528c3ff7
branch: 3.10
author: Miss Islington (bot) <31488909+miss-islington at users.noreply.github.com>
committer: miss-islington <31488909+miss-islington at users.noreply.github.com>
date: 2021-09-04T22:30:37-07:00
summary:

bpo-44571:  Add itertool recipe for a variant of takewhile() (GH-28167)

(cherry picked from commit 91be41ad933e24bff26353a19f56447e17fb6367)

Co-authored-by: Raymond Hettinger <rhettinger at users.noreply.github.com>

files:
M Doc/library/itertools.rst
M Lib/test/test_itertools.py

diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst
index fd77f99a88f577..254e055bd9eb43 100644
--- a/Doc/library/itertools.rst
+++ b/Doc/library/itertools.rst
@@ -837,6 +837,34 @@ which incur interpreter overhead.
        t1, t2 = tee(iterable)
        return filterfalse(pred, t1), filter(pred, t2)
 
+   def before_and_after(predicate, it):
+       """ Variant of takewhile() that allows complete
+           access to the remainder of the iterator.
+
+           >>> all_upper, remainder = before_and_after(str.isupper, 'ABCdEfGhI')
+           >>> str.join('', all_upper)
+           'ABC'
+           >>> str.join('', remainder)
+           'dEfGhI'
+
+           Note that the first iterator must be fully
+           consumed before the second iterator can
+           generate valid results.
+       """
+       it = iter(it)
+       transition = []
+       def true_iterator():
+           for elem in it:
+               if predicate(elem):
+                   yield elem
+               else:
+                   transition.append(elem)
+                   return
+       def remainder_iterator():
+           yield from transition
+           yield from it
+       return true_iterator(), remainder_iterator()
+
    def powerset(iterable):
        "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
        s = list(iterable)
@@ -948,4 +976,3 @@ which incur interpreter overhead.
                c, n = c*(n-r)//n, n-1
            result.append(pool[-1-n])
        return tuple(result)
-
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index a99b5e2bb71db9..6a650b9998b66d 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -2411,6 +2411,40 @@ def test_permutations_sizeof(self):
 ...             pending -= 1
 ...             nexts = cycle(islice(nexts, pending))
 
+>>> def partition(pred, iterable):
+...     "Use a predicate to partition entries into false entries and true entries"
+...     # partition(is_odd, range(10)) --> 0 2 4 6 8   and  1 3 5 7 9
+...     t1, t2 = tee(iterable)
+...     return filterfalse(pred, t1), filter(pred, t2)
+
+>>> def before_and_after(predicate, it):
+...     ''' Variant of takewhile() that allows complete
+...         access to the remainder of the iterator.
+...
+...         >>> all_upper, remainder = before_and_after(str.isupper, 'ABCdEfGhI')
+...         >>> str.join('', all_upper)
+...         'ABC'
+...         >>> str.join('', remainder)
+...         'dEfGhI'
+...
+...         Note that the first iterator must be fully
+...         consumed before the second iterator can
+...         generate valid results.
+...     '''
+...     it = iter(it)
+...     transition = []
+...     def true_iterator():
+...         for elem in it:
+...             if predicate(elem):
+...                 yield elem
+...             else:
+...                 transition.append(elem)
+...                 return
+...     def remainder_iterator():
+...         yield from transition
+...         yield from it
+...     return true_iterator(), remainder_iterator()
+
 >>> def powerset(iterable):
 ...     "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
 ...     s = list(iterable)
@@ -2538,6 +2572,21 @@ def test_permutations_sizeof(self):
 >>> list(roundrobin('abc', 'd', 'ef'))
 ['a', 'd', 'e', 'b', 'f', 'c']
 
+>>> def is_odd(x):
+...     return x % 2 == 1
+
+>>> evens, odds = partition(is_odd, range(10))
+>>> list(evens)
+[0, 2, 4, 6, 8]
+>>> list(odds)
+[1, 3, 5, 7, 9]
+
+>>> all_upper, remainder = before_and_after(str.isupper, 'ABCdEfGhI')
+>>> str.join('', all_upper)
+'ABC'
+>>> str.join('', remainder)
+'dEfGhI'
+
 >>> list(powerset([1,2,3]))
 [(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]