yield from and itertools Essentials

24 min · Intermediate

Real programs join, slice and group streams of data. The itertools module does this lazily, and yield from lets generators hand work to each other.

yield from hands work to another iterable

yield from is a statement that hands control to another iterable and yields all its values, so a generator can reuse another generator without a loop.

Looping over another iterable just to yield each item is common. yield from does the same in one line.

Python

def all_orders(online, store):    yield from online    yield from store
print(list(all_orders(["A1", "A2"], ["S1"])))

Example 1 · Joining two generators

Python

1def online():2    yield "A1"3    yield "A2"4
5def store():6    yield "S1"7
8def all_orders():9    yield from online()10    yield from store()11
12print(list(all_orders()))

Output

['A1', 'A2', 'S1']

What this code does

all_orders does no looping of its own.
Each yield from runs the other generator and passes its values on.

Each yield from passes every value of the iterable out, in order, then moves on to the next line.

yield from for nested data

Nested lists or trees need recursion. A generator can call itself with yield from and flatten any depth without building temporary lists.

Python

1def flatten(items):2    for item in items:3        if isinstance(item, list):4            yield from flatten(item)5        else:6            yield item7
8print(list(flatten([1, [2, [3, 4]], 5])))

Example 2 · Flatten any depth

Python

1def flatten(items):2    for item in items:3        if isinstance(item, list):4            yield from flatten(item)5        else:6            yield item7
8data = [1, [2, [3, [4]]], 5]9print(list(flatten(data)))

Output

[1, 2, 3, 4, 5]

What this code does

The recursion goes one level deeper for each inner list.
The result is a flat sequence in the original order.

When the item is a list, the generator delegates to a new flatten call. Otherwise it yields the value itself.

itertools: lazy building blocks

itertools is a standard-library module of fast tools that build and combine iterators lazily, such as chain, islice, groupby and combinations.

The itertools module holds fast, lazy tools that return iterators. count never ends, islice takes a part, chain joins, and accumulate keeps running totals.

Python

from itertools import accumulate, chain, count, islice
print(list(islice(count(10, 5), 3)))print(list(chain([1, 2], (3,))))print(list(accumulate([100, 50, 25])))

Example 3 · Slicing an endless stream

Python

from itertools import count, islice
evens = (n for n in count() if n % 2 == 0)print(list(islice(evens, 5)))print(list(islice(range(100), 10, 14)))

Output

[0, 2, 4, 6, 8]
[10, 11, 12, 13]

What this code does

count() goes on forever, so only islice makes the program end.
islice(it, start, stop) also accepts a start position.

islice(it, 3) is slicing for any iterator. Unlike list[:3], it works on endless streams and does not copy.

Grouping and combining

A combination picks items where order does not matter, so (A, B) and (B, A) count once. A permutation counts every order separately.

groupby collects neighbouring items with the same key, so sort first. product, permutations and combinations build every arrangement of the input.

Python

1from itertools import combinations, groupby, product2
3sales = sorted([("pen", 10), ("book", 90), ("pen", 5)])4for name, rows in groupby(sales, key=lambda r: r[0]):5    print(name, sum(v for _, v in rows))6print(list(combinations("ABC", 2)))7print(len(list(product([1, 2], "xy"))))

Example 4 · Total sales per product

Python

1from itertools import groupby2
3sales = [("pen", 10), ("book", 90), ("pen", 5), ("book", 10)]4sales.sort(key=lambda row: row[0])5for product, rows in groupby(sales, key=lambda row: row[0]):6    print(product, sum(amount for _, amount in rows))

Output

book 100
pen 15

What this code does

Sorting first puts equal keys next to each other.
Each rows is an iterator, so it must be read before moving on.

groupby only joins items that sit next to each other. Without sorting, the same key can appear in several groups.

Pairs and batches

pairwise gives each neighbouring pair. batched cuts a stream into groups of a given size. Both need Python 3.12 or newer.

Python

from itertools import batched, pairwise
temps = [20, 23, 22, 25]print([b - a for a, b in pairwise(temps)])print(list(batched(range(7), 3)))

Example 5 · Pairs, combinations and batches

Python

from itertools import batched, combinations, pairwise
print(list(pairwise([1, 4, 9, 16])))print(list(combinations(["A", "B", "C"], 2)))print(list(batched("abcdefg", 3)))

Output

[(1, 4), (4, 9), (9, 16)]
[('A', 'B'), ('A', 'C'), ('B', 'C')]
[('a', 'b', 'c'), ('d', 'e', 'f'), ('g',)]

What this code does

pairwise is a sliding window of size two.
combinations(x, 2) gives each pair once, in input order.

The last batch can be shorter. Each batch is a tuple.

Next step

Check your understanding

12 quick questions

1def flatten(items):2 for item in items:3 if isinstance(item, list):4 yield from flatten(item)5 else:6 yield item7 8data = [1, [2, [3, [4]]], 5]9print(list(flatten(data)))

1from itertools import combinations, groupby, product2 3sales = sorted([("pen", 10), ("book", 90), ("pen", 5)])4for name, rows in groupby(sales, key=lambda r: r[0]):5 print(name, sum(v for _, v in rows))6print(list(combinations("ABC", 2)))7print(len(list(product([1, 2], "xy"))))

1from itertools import groupby2 3sales = [("pen", 10), ("book", 90), ("pen", 5), ("book", 10)]4sales.sort(key=lambda row: row[0])5for product, rows in groupby(sales, key=lambda row: row[0]):6 print(product, sum(amount for _, amount in rows))