Misc minor improvements to the itertools recipes (gh-113477)

This commit is contained in:
Raymond Hettinger 2023-12-25 16:26:04 -06:00 committed by GitHub
parent 48c49739f5
commit b5dc0f83ad
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -803,11 +803,11 @@ which incur interpreter overhead.
import random import random
def take(n, iterable): def take(n, iterable):
"Return first n items of the iterable as a list" "Return first n items of the iterable as a list."
return list(islice(iterable, n)) return list(islice(iterable, n))
def prepend(value, iterable): def prepend(value, iterable):
"Prepend a single value in front of an iterable" "Prepend a single value in front of an iterable."
# prepend(1, [2, 3, 4]) --> 1 2 3 4 # prepend(1, [2, 3, 4]) --> 1 2 3 4
return chain([value], iterable) return chain([value], iterable)
@ -825,15 +825,15 @@ which incur interpreter overhead.
return starmap(func, repeat(args, times)) return starmap(func, repeat(args, times))
def flatten(list_of_lists): def flatten(list_of_lists):
"Flatten one level of nesting" "Flatten one level of nesting."
return chain.from_iterable(list_of_lists) return chain.from_iterable(list_of_lists)
def ncycles(iterable, n): def ncycles(iterable, n):
"Returns the sequence elements n times" "Returns the sequence elements n times."
return chain.from_iterable(repeat(tuple(iterable), n)) return chain.from_iterable(repeat(tuple(iterable), n))
def tail(n, iterable): def tail(n, iterable):
"Return an iterator over the last n items" "Return an iterator over the last n items."
# tail(3, 'ABCDEFG') --> E F G # tail(3, 'ABCDEFG') --> E F G
return iter(collections.deque(iterable, maxlen=n)) return iter(collections.deque(iterable, maxlen=n))
@ -848,7 +848,7 @@ which incur interpreter overhead.
next(islice(iterator, n, n), None) next(islice(iterator, n, n), None)
def nth(iterable, n, default=None): def nth(iterable, n, default=None):
"Returns the nth item or a default value" "Returns the nth item or a default value."
return next(islice(iterable, n, None), default) return next(islice(iterable, n, None), default)
def quantify(iterable, pred=bool): def quantify(iterable, pred=bool):
@ -856,7 +856,7 @@ which incur interpreter overhead.
return sum(map(pred, iterable)) return sum(map(pred, iterable))
def all_equal(iterable): def all_equal(iterable):
"Returns True if all the elements are equal to each other" "Returns True if all the elements are equal to each other."
g = groupby(iterable) g = groupby(iterable)
return next(g, True) and not next(g, False) return next(g, True) and not next(g, False)
@ -873,6 +873,30 @@ which incur interpreter overhead.
# first_true([a,b], x, f) --> a if f(a) else b if f(b) else x # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
return next(filter(pred, iterable), default) return next(filter(pred, iterable), default)
def unique_everseen(iterable, key=None):
"List unique elements, preserving order. Remember all elements ever seen."
# unique_everseen('AAAABBBCCDAABBB') --> A B C D
# unique_everseen('ABBcCAD', str.casefold) --> A B c D
seen = set()
if key is None:
for element in filterfalse(seen.__contains__, iterable):
seen.add(element)
yield element
else:
for element in iterable:
k = key(element)
if k not in seen:
seen.add(k)
yield element
def unique_justseen(iterable, key=None):
"List unique elements, preserving order. Remember only the element just seen."
# unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
# unique_justseen('ABBcCAD', str.casefold) --> A B c A D
if key is None:
return map(operator.itemgetter(0), groupby(iterable))
return map(next, map(operator.itemgetter(1), groupby(iterable, key)))
def iter_index(iterable, value, start=0, stop=None): def iter_index(iterable, value, start=0, stop=None):
"Return indices where a value occurs in a sequence or iterable." "Return indices where a value occurs in a sequence or iterable."
# iter_index('AABCADEAF', 'A') --> 0 1 4 7 # iter_index('AABCADEAF', 'A') --> 0 1 4 7
@ -893,31 +917,17 @@ which incur interpreter overhead.
except ValueError: except ValueError:
pass pass
def iter_except(func, exception, first=None): def sliding_window(iterable, n):
""" Call a function repeatedly until an exception is raised. "Collect data into overlapping fixed-length chunks or blocks."
# sliding_window('ABCDEFG', 4) --> ABCD BCDE CDEF DEFG
Converts a call-until-exception interface to an iterator interface. it = iter(iterable)
Like builtins.iter(func, sentinel) but uses an exception instead window = collections.deque(islice(it, n-1), maxlen=n)
of a sentinel to end the loop. for x in it:
window.append(x)
Examples: yield tuple(window)
iter_except(functools.partial(heappop, h), IndexError) # priority queue iterator
iter_except(d.popitem, KeyError) # non-blocking dict iterator
iter_except(d.popleft, IndexError) # non-blocking deque iterator
iter_except(q.get_nowait, Queue.Empty) # loop over a producer Queue
iter_except(s.pop, KeyError) # non-blocking set iterator
"""
try:
if first is not None:
yield first() # For database APIs needing an initial cast to db.first()
while True:
yield func()
except exception:
pass
def grouper(iterable, n, *, incomplete='fill', fillvalue=None): def grouper(iterable, n, *, incomplete='fill', fillvalue=None):
"Collect data into non-overlapping fixed-length chunks or blocks" "Collect data into non-overlapping fixed-length chunks or blocks."
# grouper('ABCDEFG', 3, fillvalue='x') --> ABC DEF Gxx # grouper('ABCDEFG', 3, fillvalue='x') --> ABC DEF Gxx
# grouper('ABCDEFG', 3, incomplete='strict') --> ABC DEF ValueError # grouper('ABCDEFG', 3, incomplete='strict') --> ABC DEF ValueError
# grouper('ABCDEFG', 3, incomplete='ignore') --> ABC DEF # grouper('ABCDEFG', 3, incomplete='ignore') --> ABC DEF
@ -932,16 +942,9 @@ which incur interpreter overhead.
case _: case _:
raise ValueError('Expected fill, strict, or ignore') raise ValueError('Expected fill, strict, or ignore')
def sliding_window(iterable, n):
# sliding_window('ABCDEFG', 4) --> ABCD BCDE CDEF DEFG
it = iter(iterable)
window = collections.deque(islice(it, n-1), maxlen=n)
for x in it:
window.append(x)
yield tuple(window)
def roundrobin(*iterables): def roundrobin(*iterables):
"roundrobin('ABC', 'D', 'EF') --> A D E B F C" "Visit input iterables in a cycle until each is exhausted."
# roundrobin('ABC', 'D', 'EF') --> A D E B F C
# Recipe credited to George Sakkis # Recipe credited to George Sakkis
num_active = len(iterables) num_active = len(iterables)
nexts = cycle(iter(it).__next__ for it in iterables) nexts = cycle(iter(it).__next__ for it in iterables)
@ -964,11 +967,43 @@ which incur interpreter overhead.
return filterfalse(pred, t1), filter(pred, t2) return filterfalse(pred, t1), filter(pred, t2)
def subslices(seq): def subslices(seq):
"Return all contiguous non-empty subslices of a sequence" "Return all contiguous non-empty subslices of a sequence."
# subslices('ABCD') --> A AB ABC ABCD B BC BCD C CD D # subslices('ABCD') --> A AB ABC ABCD B BC BCD C CD D
slices = starmap(slice, combinations(range(len(seq) + 1), 2)) slices = starmap(slice, combinations(range(len(seq) + 1), 2))
return map(operator.getitem, repeat(seq), slices) return map(operator.getitem, repeat(seq), slices)
def iter_except(func, exception, first=None):
""" Call a function repeatedly until an exception is raised.
Converts a call-until-exception interface to an iterator interface.
Like builtins.iter(func, sentinel) but uses an exception instead
of a sentinel to end the loop.
Priority queue iterator:
iter_except(functools.partial(heappop, h), IndexError)
Non-blocking dictionary iterator:
iter_except(d.popitem, KeyError)
Non-blocking deque iterator:
iter_except(d.popleft, IndexError)
Non-blocking iterator over a producer Queue:
iter_except(q.get_nowait, Queue.Empty)
Non-blocking set iterator:
iter_except(s.pop, KeyError)
"""
try:
if first is not None:
# For database APIs needing an initial call to db.first()
yield first()
while True:
yield func()
except exception:
pass
def before_and_after(predicate, it): def before_and_after(predicate, it):
""" Variant of takewhile() that allows complete """ Variant of takewhile() that allows complete
access to the remainder of the iterator. access to the remainder of the iterator.
@ -980,12 +1015,12 @@ which incur interpreter overhead.
>>> ''.join(remainder) # takewhile() would lose the 'd' >>> ''.join(remainder) # takewhile() would lose the 'd'
'dEfGhI' 'dEfGhI'
Note that the first iterator must be fully Note that the true iterator must be fully consumed
consumed before the second iterator can before the remainder iterator can generate valid results.
generate valid results.
""" """
it = iter(it) it = iter(it)
transition = [] transition = []
def true_iterator(): def true_iterator():
for elem in it: for elem in it:
if predicate(elem): if predicate(elem):
@ -993,41 +1028,8 @@ which incur interpreter overhead.
else: else:
transition.append(elem) transition.append(elem)
return return
def remainder_iterator():
yield from transition
yield from it
return true_iterator(), remainder_iterator()
def unique_everseen(iterable, key=None): return true_iterator(), chain(transition, it)
"List unique elements, preserving order. Remember all elements ever seen."
# unique_everseen('AAAABBBCCDAABBB') --> A B C D
# unique_everseen('ABBcCAD', str.lower) --> A B c D
seen = set()
if key is None:
for element in filterfalse(seen.__contains__, iterable):
seen.add(element)
yield element
# For order preserving deduplication,
# a faster but non-lazy solution is:
# yield from dict.fromkeys(iterable)
else:
for element in iterable:
k = key(element)
if k not in seen:
seen.add(k)
yield element
# For use cases that allow the last matching element to be returned,
# a faster but non-lazy solution is:
# t1, t2 = tee(iterable)
# yield from dict(zip(map(key, t1), t2)).values()
def unique_justseen(iterable, key=None):
"List unique elements, preserving order. Remember only the element just seen."
# unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
# unique_justseen('ABBcCAD', str.lower) --> A B c A D
if key is None:
return map(operator.itemgetter(0), groupby(iterable))
return map(next, map(operator.itemgetter(1), groupby(iterable, key)))
The following recipes have a more mathematical flavor: The following recipes have a more mathematical flavor:
@ -1562,16 +1564,16 @@ The following recipes have a more mathematical flavor:
>>> list(unique_everseen('AAAABBBCCDAABBB')) >>> list(unique_everseen('AAAABBBCCDAABBB'))
['A', 'B', 'C', 'D'] ['A', 'B', 'C', 'D']
>>> list(unique_everseen('ABBCcAD', str.lower)) >>> list(unique_everseen('ABBCcAD', str.casefold))
['A', 'B', 'C', 'D'] ['A', 'B', 'C', 'D']
>>> list(unique_everseen('ABBcCAD', str.lower)) >>> list(unique_everseen('ABBcCAD', str.casefold))
['A', 'B', 'c', 'D'] ['A', 'B', 'c', 'D']
>>> list(unique_justseen('AAAABBBCCDAABBB')) >>> list(unique_justseen('AAAABBBCCDAABBB'))
['A', 'B', 'C', 'D', 'A', 'B'] ['A', 'B', 'C', 'D', 'A', 'B']
>>> list(unique_justseen('ABBCcAD', str.lower)) >>> list(unique_justseen('ABBCcAD', str.casefold))
['A', 'B', 'C', 'A', 'D'] ['A', 'B', 'C', 'A', 'D']
>>> list(unique_justseen('ABBcCAD', str.lower)) >>> list(unique_justseen('ABBcCAD', str.casefold))
['A', 'B', 'c', 'A', 'D'] ['A', 'B', 'c', 'A', 'D']
>>> d = dict(a=1, b=2, c=3) >>> d = dict(a=1, b=2, c=3)