A search page shows the first ten results. The results arrive from a slow
service one page at a time, as a list per page. You want ten results without
fetching every page. Two tools from itertools do this together. The chain
function joins several iterables into one stream, and islice takes a slice
of any iterable, list or not.
from collections.abc import Iterator
from itertools import chain, islice
fetched: list[int] = []
def pages() -> Iterator[list[int]]:
for number, results in enumerate([[1, 2], [3, 4], [5, 6]]):
fetched.append(number)
yield results
first_four = list(islice(chain.from_iterable(pages()), 4))
assert first_four == [1, 2, 3, 4]
assert fetched == [0, 1] # the third page was never fetched
pages is a generator function, like the ones in "Generators and One-Shot
Iterators". chain and islice are lazy in the same way. They do work only
when something asks for the next item. Here list asks for items one at a
time. After the fourth, islice stops asking, so chain never asks for the
third page.
Pick chain or chain.from_iterable
from itertools import chain
assert list(chain([1, 2], [3])) == [1, 2, 3]
assert list(chain.from_iterable([[1, 2], [], [3]])) == [1, 2, 3]
chain(a, b, c) takes its sources as separate arguments.
chain.from_iterable(sources) takes one iterable of sources and asks for the
next source only when it needs one. Writing chain(*pages()) looks similar,
but the * runs the whole generator before chain starts, so every page is
fetched up front.
Take a slice of any iterable
from itertools import count, islice
source = iter(range(10))
assert list(islice(source, 1, 6, 2)) == [1, 3, 5]
assert next(source) == 6
assert list(islice(count(), 3)) == [0, 1, 2]
islice(iterable, stop) and islice(iterable, start, stop, step) work like
slicing a list, with two differences. Negative numbers are not allowed,
because an iterator does not know where it ends. And islice reads the items
it skips, so the source moves past them. In the example, islice read items
0 to 5, and next(source) then gets 6.
itertools.count() counts up forever. Taking three items from it with
islice is safe, because islice stops asking after the third.
islice(source, 0) reads nothing at all. For n above 0, chain still fetches
every empty page that comes before the items you want. The only way to find
out that a page is empty is to fetch it. The list you build holds at most n
items, however long the input is.
Give mixed items a type
from collections.abc import Iterator
from itertools import chain
numbers = [1, 2]
words = ["three"]
merged: Iterator[int | str] = chain(numbers, words)
lengths: list[int] = []
for item in merged:
lengths.append(item if isinstance(item, int) else len(item))
assert lengths == [1, 2, 5]
Without the annotation, mypy infers chain[object] for a chain of ints and
strings, and then rejects len(item). The annotation Iterator[int | str]
keeps both types. You still check which one each item is with isinstance
before you use something that only one of them has.
A chain is not a list
A list has a length, supports indexing, and you can loop over it many times.
A chain object is a one-shot iterator. It has no len, no indexing, and a
second loop over it is empty. Suppose a function used to return a list, and
you change it to return a chain. The first loop gives the same values, but
callers can no longer do everything they did before.
Return a list when callers need those things. Return an iterator when one pass
is all they need, and say so in the return type. Calling len on a chain
raises TypeError, and counting the items in a loop uses them up.
An Iterable parameter does not promise that you can loop over it twice,
because lists and one-shot iterators both match it.