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

Author: rhettinger

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