A Guide to using Python Itertools Module

Itertools Module

In this article, we’ll take a look at using the Python itertools Module.

This module is very useful if you want to create different types of iterators suitable for various tasks.

If you’re able to learn some of the methods of this module, this could be a very useful addition to your toolbox! Let’s get started now, by going through some useful methods.


Python Itertools – Useful Methods

In this section, we’ll look at some useful methods which generate iterators.

To use this module, we must first import it. This is already available in the standard library, so it’s pre-installed!

import itertools

Using Python itertools.chain() to chain iterables together

The Python itertools.chain() method generates an iterator from multiple iterables.

This simply chains all the iterables together into one sequence and returns a single iterator to that combined sequence.

The syntax for this method is as follows

iterator = itertools.chain(*sequence)

Let’s look at a simple example, to understand this.

import itertools

list1 = ['hello', 'from', 'AskPython']
list2 = [10, 20, 30, 40, 50]
dict1 = {'site': 'AskPython', 'url': 'https://askpython.com'}

# We can combine lists and dicts (iterables) into a single chain
for item in itertools.chain(list1, list2, dict1):
    print(item)

Here, we use the iterator directly, by iterating through it, using for item in ...

Output

hello
from
AskPython
10
20
30
40
50
site
url

Here, although we get the contents of our lists correctly, the dictionary values are not displayed.

To fix this, we could use dict.items() to get a tuple of (key, value) pairs.

import itertools

list1 = ['hello', 'from', 'AskPython']
list2 = [10, 20, 30, 40, 50]
dict1 = {'site': 'AskPython', 'url': 'https://askpython.com'}

# We can combine lists and dicts (iterables) into a single chain
for item in itertools.chain(list1, list2, dict1.items()):
    print(item)

Output

hello
from
AskPython
10
20
30
40
50
('site', 'AskPython')
('url', 'https://askpython.com')

Indeed, we now have the values printed as well, using dict1.items() as the iterable!

Using Python itertools.count() to generate a counter-based sequence

We can use the function Python itertools.count() to make iterators corresponding to a count.

iterator = itertools.count(start=0, step=1)

Here, this is an iterator which keeps counting indefinitely, from 0 onward.

This keeps increasing the count by step=1. We can also set this to a decimal/negative number.

For example, if you want to prove that you have an infinite loop, you can run the below snippet, but it is NOT recommended.

Just make sure that you can understand that itertools.count() counts infinitely.

for num in itertools.count(start=0, step=1):
    # Infinite loop!
    print(num)

Now, while you may not find the use of this function immediately obvious, you can combine it with other functions such as the zip method to construct sequences.

Consider the below example:

import itertools
numbers = [100, 200, 300, 400]

data = list(zip(itertools.count(0, 10), numbers))

print(data)

Here, you can now see the power of iterators! Since iterators produce outputs only on demand, we can zip() it with another finite iterable, such as a list!

Now, this is used to construct indices for items in the list, as you can verify using the output!

[(0, 100), (10, 200), (20, 300), (30, 400)]

Now, if you want to have a subset of the iterator sequence using Python itertools.count(), you can also use itertools.islice() to construct only a slice of the iterator.

import itertools

for num in itertools.islice(itertools.count(start=0, step=10), 4):
    print(num)

for num in itertools.islice(itertools.count(), 0, 50, 10):
    print(num)

Output

0
10
20
30
0
10
20
30
40

As you can observe, both the sequences are identical. This shows that you can have multiple approaches to generate sequences!

Use the method which you see fit, based on the problem to solve!

Using itertools.repeat() to repeat a value

Suppose you want to repeat a particular value, you can construct an iterator for the repeated value using itertools.repeat(value).

For example, if you want to construct a sequence of the form (i, 5), where i ranges from 0 to 10, you can use this function!

import itertools

data = list(zip(range(10), itertools.repeat(5)))
print(data)

Output

[(0, 5),
 (1, 5),
 (2, 5),
 (3, 5),
 (4, 5),
 (5, 5),
 (6, 5),
 (7, 5),
 (8, 5),
 (9, 5)]

Indeed, we were able to make this sequence easily!

Another example which this function is useful is if you’re trying to construct squares using map() in Python.

squares = list(map(pow, range(10), itertools.repeat(2)))
print(squares)

Output

[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

See how easily we were able to construct it using map()?

Using itertools.tee() to clone sequences

There is another useful function called tee(), which clones a sequence, and produces two sequences.

cloned1, cloned2 = itertools.tee(original_sequence)

This is based on the Linux tee command, which clones its outputs.

Here, when you clone a sequence using tee(), you cannot use the same iterator again. As a result, you must be very careful when using this function!

import itertools

single_iterator = itertools.islice(itertools.count(), 3)
cloned1, cloned2 = itertools.tee(single_iterator)

for num in cloned1:
    print('cloned1: {}'.format(num))
for num in cloned2:
    print('cloned2: {}'.format(num))

Output

cloned1: 0
cloned1: 1
cloned1: 2
cloned2: 0
cloned2: 1
cloned2: 2

Indeed, we could see two cloned sequences, having the same outputs!

Cycle through sequences using itertools.cycle()

The itertools.cycle() function provides an iterator that we can cycle through indefinitely!

This is useful if you want to keep switching between states in your application.

Consider two states of a bulb: “on” and “off”.

You can construct an iterator which cycles through the two states whenever the switch is pressed!

import itertools

# Initially, bulb is switched off, so off is the first element in the list
bulb_states = itertools.cycle(["off", "on"])

for _ in range(5):
    # Use next(iterator) to get the current state
    curr_state = next(bulb_states)
    print(f"Bulb state currently {curr_state}")

Output

Bulb state currently off
Bulb state currently on
Bulb state currently off
Bulb state currently on
Bulb state currently off

Indeed, as you can see, the bulb state keeps getting cycled between the two values “on” and “off”!

Filter items using takewhile() and dropwhile()

We can use the Python itertools.takewhile() function to filter sequence items as long as a condition is True. If the condition becomes False, it stops filtering.

iterator = itertools.takewhile(condition, *sequence)

Here is a simple example, which filters numbers, as long as the number is positive.

import itertools

sequence = itertools.takewhile(lambda x: x > 0, [1, 2, 3, -1, 10])

for item in sequence:
    print(item)

Output

1
2
3

Here, the sequence stopped after 3, since the next element is -1.

Similarly, the itertools.dropwhile() filters elements as long as a condition is False and returns all the elements after the first non-false value.

import itertools

data = itertools.dropwhile(lambda x: x < 5, [3, 12, 7, 1, -5])
for item in data:
    print(item)

Output

12
7
1
-5

Construct combinations using combinations()

We can also construct combination sequences using Python itertools.combinations().

iterator = itertools.combinations(*sequence, r)

Here is a simple example:

import itertools
words = ['hello', 'from', 'AskPython', 'how']
results = itertools.combinations(words, 2)
for item in results:
    print(item)

Output

('hello', 'from')
('hello', 'AskPython')
('hello', 'how')
('from', 'AskPython')
('from', 'how')
('AskPython', 'how')

If you want to have repetitions of consecutive elements in the combinations, you can use combinations_with_replacement().

results = itertools.combinations_with_replacement(words, 3)

for item in results:
    print(item)

Output

('hello', 'hello', 'hello')
('hello', 'hello', 'from')
('hello', 'hello', 'AskPython')
('hello', 'hello', 'how')
('hello', 'from', 'from')
('hello', 'from', 'AskPython')
('hello', 'from', 'how')
('hello', 'AskPython', 'AskPython')
('hello', 'AskPython', 'how')
('hello', 'how', 'how')
('from', 'from', 'from')
('from', 'from', 'AskPython')
('from', 'from', 'how')
('from', 'AskPython', 'AskPython')
('from', 'AskPython', 'how')
('from', 'how', 'how')
('AskPython', 'AskPython', 'AskPython')
('AskPython', 'AskPython', 'how')
('AskPython', 'how', 'how')
('how', 'how', 'how')

Similarly, you could list the permutations using permutations() and permutations_with_replacement().

That concludes some of the important functions of this module. For more functions, you can consult the official documentation.


Conclusion

In this article, we looked at various functions in the Python itertools module. Depending on your problem, you could use one of multiple approaches, to construct sequences quickly!

References