Itertools In Python

Introduction

The Python Itertools modules is a collection of functions for creating iterators that help in effective looping, every function of Itertools modules always returns an Iterator for better looping. for example, if we must iterate over a list and dictionary simultaneously within a function then two loops needs to be written one for a list and one for dictionary or combining the list elements in groups, these types of iterations and processing will be complex in nature and error-prone. Itertools module functions simplify such processing to a great extent and are extremely easy to implement.

The itertools module is considered as Gem of Python language. The article explains some of the important functions in Itertools library and a few other details about the module. Let’s explore.

Types of Iterators

There are two types of Iterators, Infinite, and finite. Infinite iterators continue to run forever if no stopping condition is encountered. Finite iterators terminate. Itertools module has various functions of each type

  1. Finite Iterator functions: chain, compress, tee, dropWhile, takeWhile
  2. Infinite Iterator functions: count, repeat, cycle

Finite Iterators

chain function: The chain function can take any number of Iterables and chain all the iterables into one single sequence and returns a single iterator.

from itertools import *
names = ["John", "David", "James"]
age = [40, 50, 60]
for element in chain(names, age):
    print(element)

# John
# David
# James
# 40
# 50
# 60

The ‘chain’ function takes two list objects, but it is capable of taking all types of iterables, we can modify the above code to take a dictionary object as well.

names = ["John", "David", "James"]
age = [40, 50, 60]
dict1 = {"state":"TX"}
for element in chain(names, age, dict1.items()):
    print(element)


#John
#David
#James
#40
#50
#60
#('state', 'TX')

compress function: The function takes an iterable and selector, then it filters the data that have the value in the selector to True and returns an iterator.

names = ['John', 'David', 'James']
select = [False, True, True]
for name in compress(names, select):
    print(name)

# David
# James

Eliminates the elements with selector value as False.

tee: The tee function takes an iterable, clones it, and returns a new iterator. Additionally, the tee function can also accept a number (n) and clone the iterable ‘n’ times.  

numbers = [1,2,3]
i, j = itertools.tee(numbers)
print(list(i), list(j))

# [1, 2, 3] [1, 2, 3]
numbers = [1,2,3]
it = itertools.tee(numbers, 5)
for i in it:
    print(list(i))

#[1, 2, 3]
#[1, 2, 3]
#[1, 2, 3]
#[1, 2, 3]
#[1, 2, 3]

dropWhile: The dropWhile function takes and iterable and filters the elements which satisfy the given condition and returns an iterator.

ages = [0,0,10,20,30]
for a in dropwhile(lambda x : x == 0, ages):
 print(a)

#10
#20
#30

takeWhile: The takeWhile function takes an iterable and filters the elements as long as the condition is satisfied.

ages = [10,20,0,30]
result = itertools.takewhile(lambda x: x > 0, ages)
for a in result:
    print(a)

#10
#20

Infinite Iterators

count: The count function takes the starting value like 0, 1, or 2 indicating the start value and another step parameter which is for incrementing the counter value at every step.

iterator = itertools.count(start, step)

for i in itertools.count(0, 1):
    if(i < 5):
        print(i)
    else:
        break  

repeat: The repeat function takes an iterable and an optional times parameter, the times parameter is a terminating condition, if this parameter is not provided the repeat method will keep running forever.

it = itertools.repeat('Hello World', times = 5)
for w in it:
    print (w)

#Hello World
#Hello World
#Hello World
#Hello World
#Hello World

cycle: The cycle function is the same as repeat it keeps running forever and we have to explicitly provide a terminating condition.

l = ['a', 'b', 'c']
i = 0
for e in itertools.cycle(l):
    if(i < 5):
        print(e)
    else:
        break
    
    i = i +1

Summary

The article covered types of iterators and a few important Itertools modules functions, the simplicity of Itertools modules is amazing it helps in ensuring that code is less error-prone, clean as most of the complexity is abstracted out by the module.