Building Pipelines with itertools
Ready-made parts for working with sequences
In this lesson
Explain it like I’m 5
itertools is a drawer of pre-made pipe fittings. You are connecting parts, not machining your own.
The ones you will actually use
The module has around twenty functions. These five cover nearly everything a beginner needs:
chain(a, b), treat several sequences as one continuous run.islice(source, n), take the first n items, including from something endless.count(start), an endless counter;cycle(items), repeat forever.groupby(items, key), group neighboring items that share a key.zip_longest(a, b), pair up sequences of different lengths.
from itertools import chain, islice, count
first = ["a", "b", "c"]
second = ["d", "e", "f"]
print(list(chain(first, second)))
print(list(islice(chain(first, second), 4)))
# islice is what makes an endless source safe to use:
print(list(islice(count(10), 5)))
['a', 'b', 'c', 'd', 'e', 'f'] ['a', 'b', 'c', 'd'] [10, 11, 12, 13, 14]
groupby needs sorted data
groupby collects neighboring items that share a key. It does not search the whole sequence, which makes it fast and lazy, and means unsorted input produces fragmented groups.
So the rule is: sort by the same key first. This catches everyone once.
from itertools import groupby
books = [
{"genre": "scifi", "title": "Dune"},
{"genre": "crime", "title": "Gorky Park"},
{"genre": "scifi", "title": "Neuromancer"},
]
get_genre = lambda book: book["genre"]
print("unsorted:")
for genre, group in groupby(books, key=get_genre):
print(" ", genre, [b["title"] for b in group])
print("sorted first:")
for genre, group in groupby(sorted(books, key=get_genre), key=get_genre):
print(" ", genre, [b["title"] for b in group])
unsorted: scifi ['Dune'] crime ['Gorky Park'] scifi ['Neuromancer'] sorted first: crime ['Gorky Park'] scifi ['Dune', 'Neuromancer']
Chaining parts into a pipeline
Because every one of these returns an iterator, you can feed one straight into the next. Nothing is computed until something actually consumes the result, and even then only one item at a time.
The practical effect: a pipeline over a million-row file uses the same memory as one over ten rows.
from itertools import chain, islice
batch_one = range(1, 5)
batch_two = range(100, 105)
# Nothing runs yet; this just describes the work.
pipeline = (n * 2 for n in chain(batch_one, batch_two) if n % 2 == 0)
print(list(islice(pipeline, 3)))
[4, 8, 200]
Combine two lists and take the first five items, without building a merged list. Use chain to join them and islice to take five.
from itertools import chain, islice
first = ["a", "b", "c"]
second = ["d", "e", "f"]
# TODO: take the first five items across both lists
for item in []:
print(item)Nest them: islice(chain(first, second), 5). chain makes the two lists look like one sequence, and islice stops after five.
from itertools import chain, islice
first = ["a", "b", "c"]
second = ["d", "e", "f"]
for item in islice(chain(first, second), 5):
print(item)Common mistake: Calling groupby on unsorted data
It looks like it should find every matching item wherever it is.
It only groups neighbors. Sort by the same key first, and use the same key function for both.
Common mistake: Wrapping the pipeline in list() too early
list() makes it printable, so it gets added while debugging and left behind.
That materializes everything and throws away the memory benefit. Keep it lazy until the final step.
Common mistake: Iterating count() or cycle() with no limit
They look like ordinary sequences.
They never end. Always pair them with islice or a break.
What must be true of your data before using groupby?
groupby only groups neighboring items, so sorting by that key first is what makes the grouping complete.
What does islice(count(10), 5) produce?
count(10) counts upward forever; islice takes the first five, which is what makes it safe.
Why is a chained pipeline memory-efficient?
Every stage is an iterator, so items flow through one at a time and nothing intermediate is stored.
Mini exercise (medium)
You have two lists of log lines. Write first_errors(a, b, limit) that treats them as one sequence, keeps only the lines containing "ERROR", and returns at most limit of them, without building a combined list.
Give it a shot. Complete the code and press Run; there’s nothing to download or configure.
from itertools import chain, islice
batch_a = ["INFO ok", "ERROR disk full", "INFO done"]
batch_b = ["ERROR timeout", "INFO idle", "ERROR refused"]
def first_errors(a, b, limit):
return [] # TODO: chain the sources, keep ERROR lines, cap at limit
print(first_errors(batch_a, batch_b, 2))
print(first_errors(batch_a, batch_b, 10))chain(a, b) joins the sources. Filter with a generator expression, then cap the result with islice(..., limit) and wrap the final answer in list().
from itertools import chain, islice
batch_a = ["INFO ok", "ERROR disk full", "INFO done"]
batch_b = ["ERROR timeout", "INFO idle", "ERROR refused"]
def first_errors(a, b, limit):
errors = (line for line in chain(a, b) if "ERROR" in line)
return list(islice(errors, limit))
print(first_errors(batch_a, batch_b, 2))
print(first_errors(batch_a, batch_b, 10))