Mastering Python Standard Library: itertools.chain
Vladislav Zenin
Posted on December 10, 2022
Imagine, you need to iterate over some N iterables.
For example, you have two lists: l1 and l2.
In [2]: l1 = list(range(5))
In [3]: l2 = list(range(10))
In [4]: l1
Out[4]: [0, 1, 2, 3, 4]
In [5]: l2
Out[5]: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
Here is the easiest way to do so:
for i in l1+l2: print(i, end=", ")
# 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
However, it may not be the best one. l1+l2
statement is a list concatenation, and that give you a new list with len(l1+l2) == len(l1) + len(l2)
. If you positive that both lists are rather small, then it's kinda okay.
But, let us assume they are each of 1GB in RAM. At peak, your program will consume 4GB, twice the size of input lists. And what if you don't have much RAM? - maybe your code is in AWS Lambda, etc.
Actually, we want to do something like this:
def gen(l1, l2):
yield from l1
yield from l2
for i in gen(l1,l2): print(i, end=", ")
# 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
No new lists, no copies, no memory overhead. Just iterate over the first list and then iterate over the second one.
And that gen
iterator is already coded for you, and also known as itertools.chain
import itertools
for i in itertools.chain(l1,l2): print(i, end=", ")
# 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
By the way, there is another form of itertools.chain
, itertools.chain.from_iterable
. It does absolutely the same, except input arguments unpacking:
for i in itertools.chain.from_iterable([l1, l2]): print(i, end=", ")
# 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
So, in general:
# this is itertools.chain
def my_chain(*collections):
for collection in collections:
yield from collection
# this is itertools.chain.from_iterable
def my_chain_from_iterable(collections):
for collection in collections:
yield from collection
Why there are 2 chains, with one tiny "*" difference? I really don't know - but who am I to judge authors of itertools module, they are true gods.
But I do know, that "entities should not be multiplied beyond necessity". And this thought brings us back to our unnecessary extra list creation issue.
So what’s the point?
Well, use chain! Learn itertools
module. Think about performance. Save the memory, in production environment it is actually limited and not really cheap!
Anything else to read?
Sure.
Whole lotta docs - Master the power of standard library!
Itertools module docs - chain is not the only one, there are plenty more
Occam's Razor - really, read it
Posted on December 10, 2022
Join Our Newsletter. No Spam, Only the good stuff.
Sign up to receive the latest update from our blog.