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

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)
-

Lib/test/test_itertools.py

@@ -2412,6 +2412,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)
@@ -2539,6 +2573,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)]