bpo-44571: Add itertool recipe for a variant of takewhile() (GH-28167)
(cherry picked from commit 91be41ad933e24bff26353a19f56447e17fb6367)
Co-authored-by: Raymond Hettinger <rhettinger@users.noreply.github.com>
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index a99b5e2..6a650b9 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)]