[three]Bean

Using itertools to un-nest your code

Mar 15, 2011 | categories: python View Comments

Do you ever write code that becomes way nested and totally out of hand? Being smart with itertools can help un-indent your biz.

Below follow two functions definitions: nested_generator and itertools_generator. Put them in a module named itertools_blogpost.py.

from itertools import product

def nested_generator(carnivores, herbivores, plants):
    """ You do this every day. """
    for carnivore in carnivores:
        for herbivore in herbivores:
            for plant in plants:
                yield "%s eat %s eat %s" % (carnivore, herbivore, plant)

def itertools_generator(carnivores, herbivores, plants):
    """ This can 'flatten' your code,
    make it less nested, more pythonic. """
    for carnivore, herbivore, plant in product(carnivores, herbivores, plants):
        yield "%s eat %s eat %s" % (carnivore, herbivore, plant)

Let's test it with another script.

#!/usr/bin/env python

from itertools_blogpost import nested_generator, itertools_generator

if __name__ == '__main__':
    # Operate on these lists
    carnivores = ["lions", "tigers", "bears"]
    herbivores = ["hippies"]
    plants = ["carrots", "beets", "garlic"]

    # Test to show that the two functions are `equivalent`
    result1 = list(nested_generator(carnivores, herbivores, plants))
    result2 = list(itertools_generator(carnivores, herbivores, plants))
    are_they_equal = (result1 == result2)
    print "Gravey? -->", are_they_equal # Gravey!

    # Test to show that they run in an equivalent amount of time
    from timeit import Timer
    nested_timer = Timer(
        """
        carnivores = ["lions", "tigers", "bears"]
        herbivores = ["hippies"]
        plants = ["carrots", "beets", "garlic"]
        nested_generator(carnivores, herbivores, plants)
        """, "from itertools_blogpost import nested_generator")

    itertools_timer = Timer(
        """
        carnivores = ["lions", "tigers", "bears"]
        herbivores = ["hippies"]
        plants = ["carrots", "beets", "garlic"]
        itertools_generator(carnivores, herbivores, plants)
        """, "from itertools_blogpost import itertools_generator")

    nested_times = nested_timer.repeat()
    itertools_times =  itertools_timer.repeat()
    n = len(nested_times)

    print "          Nested  Iter"
    print "Maximum   {0:2.3f}   {1:2.3f}".format(
        max(nested_times), max(itertools_times))
    print "Minimum   {0:2.3f}   {1:2.3f}".format(
        min(nested_times), min(itertools_times))
    print "Average   {0:2.3f}   {1:2.3f}".format(
        sum(nested_times)/n, sum(itertools_times)/n)

Their output is equivalent and (run it for yourself) the times are roughly equivalent. Sometimes one comes out on top of the other, sometimes vice versa. No statistically significant difference.

The upshot: with itertools.product you can say goodbye to awfully nested for loops.

View Comments