You have five tasks and three workers, and you want to hand the tasks out in
turn. The itertools.cycle function repeats the workers for as long as you
need:
from itertools import cycle
workers = cycle(["ana", "ben", "cy"])
tasks = ["t1", "t2", "t3", "t4", "t5"]
print(list(zip(tasks, workers)))
# [('t1', 'ana'), ('t2', 'ben'), ('t3', 'cy'), ('t4', 'ana'), ('t5', 'ben')]
A cycle never ends on its own, so here zip stops it when tasks runs out.
Repeat items with cycle and repeat
A cycle saves each item the first time it reads it, then repeats the saved
items forever. Never call list() on one, because it never finishes.
The repeat(obj, n) function yields the same object n times, or forever if
you leave out n. It is handy with map when one argument stays the same:
from itertools import repeat
print(list(map(pow, [2, 3, 4], repeat(2)))) # [4, 9, 16]
Watch out for mutable objects. The call repeat([], 3) gives one list three
times, not three lists, so appending to one changes them all. For separate
lists, write [[] for _ in range(3)].
Handle inputs of different lengths
Plain zip stops at the shortest input and drops the rest without a warning.
Two other options let you choose:
from itertools import zip_longest
names = ["a", "b"]
scores = [1]
print(list(zip(names, scores))) # [('a', 1)]
print(list(zip_longest(names, scores))) # [('a', 1), ('b', None)]
try:
list(zip(names, scores, strict=True))
except ValueError as error:
print(error) # zip() argument 2 is shorter than argument 1
Use zip to drop the extras, and strict=True when a mismatch is a bug. Use
zip_longest to pad the short input with fillvalue, which is None unless
you pass another value. If None is real data, pass a different one.
A strict zip finds the mismatch only when the shorter input runs out, so your
loop has already handled the earlier pairs by then.
Take or skip the start of a stream
Both tools here take a predicate, a function that returns true or false for
each item. The takewhile function yields items while the predicate is true
and stops at the first item where it is false. The dropwhile function skips
items while the predicate is true, then yields the first false item and
everything after it, without testing again.
from itertools import dropwhile, takewhile
source = iter([1, 2, 0, 3])
print(list(takewhile(lambda n: n > 0, source))) # [1, 2]
print(list(source)) # [3]
print(list(dropwhile(lambda n: n > 0, [1, 2, 0, 3]))) # [0, 3]
The 0 is in neither of the first two lists. The takewhile function had to
read it to find the end of the run, and it throws that item away.
Check your understanding
You use takewhile to read a file's header lines, which start with #, then loop over the rest. The first data line is missing. Why, and how do you keep it?
Show answer
Hide answer
The takewhile function read that line to learn that the header had ended,
and it does not give the line back. Write a plain loop instead. Read lines
until one does not start with #, handle that line, then loop over the rest.
If the line that ends the header is a blank separator you do not need,
takewhile is fine.
Keep the items a test rejects
The filterfalse function keeps the items where the predicate is false, the
opposite of filter. With None as the predicate, it keeps the falsy items,
the ones that count as false in an if:
from itertools import filterfalse
print(list(filterfalse(None, [0, 1, "", "ok", None]))) # [0, '', None]
Do not split one iterator into good and bad items with filter and
filterfalse. Both read the same iterator, so the first one you loop over uses
it up. Write one loop that appends each item to one of two lists.
Call a function with each tuple
The starmap(func, rows) function calls func(*row) for each tuple in rows.
Plain map takes one item from each of several inputs instead:
from itertools import starmap
print(list(starmap(pow, [(2, 5), (3, 2)]))) # [32, 9]
print(list(map(pow, [2, 3], [5, 2]))) # [32, 9]
Use starmap when your data already comes in tuples, such as rows from a CSV
file or the pairs that zip makes.
Select items with a list of flags
The compress(data, selectors) function keeps each item whose matching
selector is true:
from itertools import compress
print(list(compress(["a", "b", "c"], [True, False, True]))) # ['a', 'c']
print(list(compress(["a", "b", "c"], [True, False]))) # ['a']
It stops without an error when either input runs out, so the second call drops
"c".
Check your understanding
A teammate cannot follow your one-line pipeline of starmap, compress and takewhile. Does splitting it into named steps or a loop use more memory?
Show answer
Hide answer
No. Each named step can be an iterator or a generator, and nothing runs until
you loop over the last one. A loop with plain if statements is often the
clearest way to show branching and error handling. Memory grows only where a
tool keeps items, such as cycle, tee or a call to list().
Going further
The itertools documentation ends with recipes built from these tools.