10 Python Itertools To Make Your Code Neater, Cleaner, and Better

Implement the same thing with shorter code

The beauty of Python lies in its simplicity.

Not only because Python’s syntax is elegant, but also due to it has many well-designed built-in modules that help us implement common functionalities efficiently.

The itertools module, which is a good example, provides many powerful tools for us to manipulate Python iterables in shorter code.

Do more by less. This is what you can get from the itertools module. Let’s check it out from this article.

1. itertools.product(): A Tricky Way To Avoid Nested Loops

When a program becomes more and more complicated, you may need to write nested loops. At the same time, your Python code will become ugly and unreadable:

list_a = [1, 2020, 70]
list_b = [2, 4, 7, 2000]
list_c = [3, 70, 7]

for a in list_a:
    for b in list_b:
        for c in list_c:
            if a + b + c == 2077:
                print(a, b, c)
# 70 2000 7

How to make the above code Pythonic again?

The itertools.product() function is your friend:

from itertools import product

list_a = [1, 2020, 70]
list_b = [2, 4, 7, 2000]
list_c = [3, 70, 7]

for a, b, c in product(list_a, list_b, list_c):
    if a + b + c == 2077:
        print(a, b, c)
# 70 2000 7

As demonstrated above, it returns the Cartesian product of input iterables, which helps us put 3 nested for loops into one.

2. itertools.compress(): A Convenient Way To Filter Data

We can filter a list of items through one or a few loops for sure.

But sometimes, we may not need to write any loops. Cause there is a function named itertools.compress().

The itertools.compress() function returns an iterator that filters an iterable based on the values of a corresponding boolean mask.

For example, the following code selects the true leaders using the itertools.compress() function:

import itertools
leaders = ['Yang', 'Elon', 'Tim', 'Tom', 'Mark']
selector = [1, 1, 0, 0, 0]
print(list(itertools.compress(leaders, selector)))
# ['Yang', 'Elon']

The second argument, selector, works as a mask, we can also define it as follows:

selector = [True, True, False, False, False]

3. itertools.groupby(): Group Elements of an Iterable

The itertools.groupby() function is a convenient way to group adjacent duplicate items of an iterable.

For instance, we can group a long string as follows:

from itertools import groupby

for key, group in groupby('YAaANNGGG'):
    print(key, list(group))
# Y ['Y']
# A ['A']
# a ['a']
# A ['A']
# N ['N', 'N']
# G ['G', 'G', 'G']

Furthermore, we can harness its second argument to tell the groupby() function how to determine whether two items are the same or not:

from itertools import groupby

for key, group in groupby('YAaANNGGG', lambda x: x.upper()):
    print(key, list(group))
# Y ['Y']
# A ['A', 'a', 'A']
# N ['N', 'N']
# G ['G', 'G', 'G']

4. itertools.combinations(): Get All Combinations of a Given Length From an Iterable

It needs a beginner a while to write a bug-free function to get all possible combinations of a list.

In fact, if she knows the itertools.combination() function, she can get it easily:

import itertools

author = ['Y', 'a', 'n', 'g']

result = itertools.combinations(author, 2)

for x in result:
    print(x)
# ('Y', 'a')
# ('Y', 'n')
# ('Y', 'g')
# ('a', 'n')
# ('a', 'g')
# ('n', 'g')

As the above program shows, the itertools.combination() function has two parameters, one is the original iterable, and the other is the length of the subsequences that will be generated by the function.

5. itertools.permutations(): Get All Permutations of a Given Length From an Iterable

Since there is a function to get all combinations, of course, there is another function named itertools.permutations to get all possible permutations:

import itertools

author = ['Y', 'a', 'n', 'g']

result = itertools.permutations(author, 2)

for x in result:
    print(x)

# ('Y', 'a')
# ('Y', 'n')
# ('Y', 'g')
# ('a', 'Y')
# ('a', 'n')
# ('a', 'g')
# ('n', 'Y')
# ('n', 'a')
# ('n', 'g')
# ('g', 'Y')
# ('g', 'a')
# ('g', 'n')

As shown above, the use of itertools.permutations() function is similar as the itertools.combinations(). The only difference is their results.

6. itertools.accumulate(): Generate Accumulated Items from an Iterable

Getting a series of accumulated values based on an iterable is a common requirement. With the help of the itertools.accumulate() function, we don’t need to write any loops to implement.

import itertools
import operator

nums = [1, 2, 3, 4, 5]
print(list(itertools.accumulate(nums, operator.mul)))
# [1, 2, 6, 24, 120]

The above program is the same as follows if we wouldn't like to use the operator.mul:

import itertools

nums = [1, 2, 3, 4, 5]
print(list(itertools.accumulate(nums, lambda a, b: a * b)))
# [1, 2, 6, 24, 120]

7. itertools.repeat(), itertools.cycle(), itertools.count(): Make Infinite Iterables

In some cases, we need to get an infinite iterable. There are 3 functions that are helpful:

itertools.repeat(): Generate the same item repeatedly

For example, we can get three same “Yang” as follows:

import itertools
print(list(itertools.repeat('Yang', 3)))
# ['Yang', 'Yang', 'Yang']

itertools.cycle(): Get an infinite iterator by cycling

The itertools.cycle function will not stop until you break the loop:

import itertools

count = 0

for c in itertools.cycle('Yang'):
    if count >= 12:
        break
    else:
        print(c, end=',')
        count += 1
# Y,a,n,g,Y,a,n,g,Y,a,n,g,

itertools.count(): generate an infinite sequence of numbers

If all we need is numbers, use the itertools.count function:

import itertools

for i in itertools.count(0, 2):
    if i == 20:
        break
    else:
        print(i, end=" ")
# 0 2 4 6 8 10 12 14 16 18

As illustrated above, its first parameter is the starting number, and the second parameter is the step.

8. itertools.pairwise(): Get Tuples of Pairs Easily

Since Python 3.10, the itertools module has a new function named pairwise. It is a small and neat tool to generate successive overlapping pairs from an iterable.

import itertools

letters = ['a', 'b', 'c', 'd', 'e']

result = itertools.pairwise(letters)

print(list(result))
# [('a', 'b'), ('b', 'c'), ('c', 'd'), ('d', 'e')]

9. itertools.takewhile(): Filter Elements in a Different Way

itertools.takewhile() returns an iterator that generates elements from an iterable as long as a given predicate function evaluates to True.

import itertools

nums = [1, 61, 7, 9, 2077]

print(list(itertools.takewhile(lambda x: x < 100, nums)))
# [1, 61, 7, 9]

This function is different from the built-in filter() function.

The filter function will go through the whole list:

nums = [1, 61, 7, 9, 2077]

print(list(filter(lambda x: x < 10, nums)))
# [1, 7, 9]

However, the itertools.takewhile function, as its name implies, will stop when the evaluating function is False:

import itertools

nums = [1, 61, 7, 9, 2077]

print(list(itertools.takewhile(lambda x: x < 10, nums)))
# [1]

10. itertools.dropwhile(): A Reverse Operation of itertools.takewhile()

This function seems like a reverse idea of the previous one.

The itertools.takewhile() returns the elements of the iterable as long as the predicate function is True, whereas itertools.dropwhile() drops the elements of the iterable as long as the predicate function is True and then returns the remaining elements.

import itertools

nums = [1, 61, 7, 9, 2077]

print(list(itertools.dropwhile(lambda x: x < 100, nums)))
# [2077]

Thanks for reading. ❤️

If you like it, please follow me to enjoy more articles about programming and technologies.

Reference: