Python's itertools

In Python one can do the same thing in a number of ways.

There are times when you don't want to use the good old for loop. Use Python's built-in itertools.

The itertools module contains functions that return iterators.

Let us apply functions from this module to something familiar like displaying a pyramid of stars and associating the English alphabet to the ASCII codes.

The examples are intentionally simple to make it easy to follow and apply later to different situations.

Create an iterator

itertools.count()

One can create a new iterator by using the count() method. The syntax is

itertools.count(start, step)

When applied to numbers, this function returns a stream of integers. The numbers begin from the start parameter and are uniformly separated by the step parameter.

The count() function returns an infinite series, so we need a way to terminate it. We will use takewhile() from itertools: it takes a predicate to break the iterable passed to it.

import itertools  
from itertools import  takewhile

# count
#takewhile(predicate, iterable)

sequence = itertools.count()
digits = list(takewhile(lambda x : x < 10, sequence))
print(digits)
print()

We have used a lambda function for the predicate in order to break the series generated by the count() function. But the predicate function can be a normal function too.

Let us map the English alphabet to the ASCII codes. We will use the count function again, but this time we will pass a start parameter.

def list_dict(sequence, start):
    assoc_array = dict(zip(sequence, itertools.count(start)))
    return assoc_array

en_alphabet = list('abcdefghijklmnopqrstuvwxyz')

ascii_dict = list_dict(en_alphabet, 97)
print(ascii_dict)
print()

itertools.cycle()

The syntax is

itertools.cycle(iter)

The cycle() method accepts an iterator and returns a copy on an infinite series. Again, in order to use it, we need to tell it where to stop. We use the range function to get a list of ASCII codes and then zip it with the list of English letters.

We have built this dictionary in a previous post with lists alone, but this time we have used itertools.

print(dict(zip(en_alphabet, itertools.cycle(range(97, 123)))))
print()

Itertools.repeat()

This function accepts an element and returns it repeatedly, unless you tell it to stop with the second parameter. Here's the syntax -

itertools.repeat(elem, [n])

We will display a simple pyramid of stars using this method. We pass a star character to it and have it repeated n times. We use the join() function in the string class to concatenate the stars.

import itertools

# pyramid of stars 
def triangle(ch, n):
    if n <= 0:
        return
    s = ''.join(itertools.repeat(ch, n))
    print(s)
    n = n - 1
    triangle(ch, n)

triangle('*', 10)

We have used recursion to display one star less in each line. Try it.

The sharp reader would have noticed that we have not used loops are all!

Hope you found this post useful. If you need more stuff like this one, examples of using modules, then reach me at anand_b@twitter.com.

Share, tweet and comment if you like.

Happy coding!

Python's itertools

count, repeat and cycle

Create an iterator