Split an Iterator with tee and Combine Inputs with product

A generator gives you each item once. As "Generators and One-Shot Iterators" showed, a second loop over it finds nothing. Sometimes two parts of your code need to read the same stream. The itertools.tee function splits one iterator into two independent iterators:

from itertools import tee

readings = iter([3, 5, 4])
left, right = tee(readings)
print(next(left), next(left))   # 3 5
print(next(right))              # 3
print(list(left), list(right))  # [4] [5, 4]

This article calls each iterator that tee returns a branch. A branch keeps its own position, and every branch sees every item.

The slowest branch decides how much tee stores

The tee function reads each item from the source once. It keeps the item in a buffer until every branch has read it. In the snippet, after left reads 3 and 5, the buffer holds both for right. As right catches up, tee drops them. Keep the branches close together, and the buffer stays small.

Common mistakes

Reading one branch to the end before the other. Then tee stores every item, because the other branch has read none of them yet. If you plan to use one branch fully and then the other, call list() on the source once and loop over the list twice.

Using the source after you split it. An item you pull from the source yourself never reaches the branches:

source = iter([1, 2, 3])
a, b = tee(source)
next(source)             # 1 goes to neither branch
print(list(a), list(b))  # [2, 3] [2, 3]

After calling tee, read only through the branches.

Changing an item through one branch. The branches share the same objects, not copies, so a list you append to through left has the new item in right too.

The product function reads all its inputs first

The itertools.product function gives every combination of one item from each input, in place of nested loops:

from itertools import product

for size, colour in product(["S", "M"], ["red", "blue"]):
    print(size, colour)
# S red
# S blue
# M red
# M blue

To make every combination, product goes over each input many times. An iterator can be read only once, so as soon as you call product, it copies each input into a tuple. Two things follow.

First, an endless input such as itertools.count() never finishes copying, so product never returns and its memory keeps growing.

Second, passing the same iterator twice gives nothing, because copying the first argument uses it up:

stream = iter([1, 2, 3])
print(list(product(stream, stream)))  # []

stream = iter([1, 2, 3])
print(len(list(product(stream, repeat=2))))  # 9

The repeat argument copies the input once and reuses the copy. Passing the same list twice also works, because each loop over a list starts again.

Count the combinations before you loop

The number of tuples is the input lengths multiplied together. Fifty values in each of 3 positions give 50 ** 3, which is 125,000 tuples, and 5 positions give 312,500,000. The tuples arrive one at a time, so memory stays small, but your loop still runs 312.5 million times. Work out the count first.

If you want each pair of different items once, use combinations instead. For n items, combinations(items, 2) gives n * (n - 1) / 2 pairs. In contrast, product(items, repeat=2) gives n * n, because it pairs each item with itself and makes every pair in both orders:

from itertools import combinations

print(list(combinations(["a", "b", "c"], 2)))
# [('a', 'b'), ('a', 'c'), ('b', 'c')]

Which itertools store data

product, combinations and permutations copy their inputs when you call them.
tee stores the items its slowest branch has not read yet.
cycle, from the next article, stores each item so it can repeat them.
chain, islice, takewhile and pairwise keep at most one item at a time.

Check your understanding You count a million log lines through one tee branch, then find the longest line through the other. How many lines does tee hold at its peak? Show answer

All million. When the first loop ends, the second branch has read nothing, so tee must keep every line for it. A list would hold the same lines with less code. Better still, count the lines and track the longest one in a single loop, and nothing needs storing.

Going further

The branches of one tee are not safe to use from several threads at once. Doing so can raise RuntimeError, even when the source itself is safe. The concurrency track covers threads.

The itertools documentation lists every tool, with recipes that combine them.