Loop Details

[ichabod801]

Our last example builds a list of values. This is another very common use of loops. Once again, Python provides us with a simpler way to do it: list comprehensions. These are not functions, like enumerate, zip, and itertools.product, but syntax for building lists using loop statements. The classic loop form of building a list is this:

some_list = []
for value in another_list:
    some_list.append(some_function(value))

The list comprehension form of this is just one line:

some_list = [some_function(value) for value in another_list]

So first we take whatever we are appending, and we put that at the beginning of the list comprehension. Then we take the for statement, and we put that at the end of the list comprehension. Note that we drop the colon from the for statement, since we are not starting a new block of code.

So our deck of cards example could become a list comprehension this way:

from itertools import product
deck = [rank + suit for rank, suit in product('A23456789TJQK', 'CDHS')]

That's the one loop version of the deck example, but you can also have nested loops in a list comprehension. Here's how you would do the nested loop version of the deck example:

deck = [rank + suit for rank in 'A23456789TJQK' for suit in 'CSHS']

Note that this gives you the exact same sequence of cards that our earlier example gave. That means that the inner loop in a list comprehension is the one that is listed second.

But wait! There's more! Often you have a conditional in your loop. Take generating a list of primes:

def is_prime(number):
    for factor in range(2, number // 2):
        if not number % factor:
            return False
    return True 

primes = []
for number in range(2, 18):
    if is_prime(number):
        primes.append(number)

Yeah, it could be more efficient, but this is just a toy example. The real question is: how do we put this in a list comprehension?

def is_prime(number):
    for factor in range(2, number // 2):
        if not number % factor:
            return False
    return True 

primes = [number for number in range(2, 18) if is_prime(number)]

As you can see, we can put the conditional in the list expression as well, right after the for statement. That's nice, but take a look at that is_prime function. It looks a lot like the loops we've been converting into list comprehensions, it's just not appending the values to a list. But if did make a list of those values, we could put it into the all() function, which returns True if all the values in it's list are True. That would turn our is_prime function into a list comprehension:

def is_prime(number):
    return all([number % factor for factor in range(2, number // 2)])

primes = [number for number in range(2, 18) if is_prime(number)]

While I may have pooh-poohed efficiency earlier with this toy example, this does make it even worse. Our old function returned False after finding the first factor of the number. This function checks all of the potential factors before returning a value. It turns out there is a subtle way to fix this:

def is_prime(number):
    return all(number % factor for factor in range(2, number // 2))

primes = [number for number in range(2, 18) if is_prime(number)]

Do you see the difference? All we did was take out the brackets for the list comprehension in the function. Now it's not a list comprehension, it's a generator. What is a generator? They're what we've been using all along: range, enumerate, zip, dict.values(), dict.items() and product are all generators. The key point with these generators is that give out their values one at a time, rather than creating all of the values at once and returning a list of them.

This is of use here because the all() function is smart. It looks at the values one at a time, and if it finds a False value, it stops looking and returns False. Just like the original version of our is_prime function did.

So before we dealt with two loops at the same time, nesting one within the other. Can we do that with list comprehensions, or with a list comprehension and a generator? It turns out we can. Here's the prime number generator in one line:

primes = [number for number in range(2, 18) if all(number % factor for factor in range(2, number // 2))]

While this may seem cool, it's actually a pet peeve I have with list comprehensions. Some people seem to think that every list than can be built with a list comprehension should be built with a list comprehension. Then end up with huge list comprehensions that are just a confusing hot mess. My advice is that if your list comprehension doesn't fit on one line, put it back in a standard loop structure. The indentation structure of a standard loop will make your code more clear to whoever is trying to understand it.

Note that lists don't have all the fun. Dictionaries have comprehensions as well. The difference is, of course, that you need two values each time through the "loop," a key and a value. For example, take this loop:

rot13 = {}
letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
for plain, cipher in zip(letters, letters[13:] + letters[:13]):
    rot13[plain] = cipher

With a list comprehension, we started with the part that was being appended. With a dictionary comprehension we start with the key and value of the assignment. But instead of an assignment (which is can't be in a list comprehension), we use a colon as in a dictionary literal. As before, the for statement goes in the second part of the comprehension. So the dictionary comprehension for the above loop is:

letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
rot13 = {plain: cipher for plain, cipher in zip(letters, letters[13:] + letters[:13])}

As with the list comprehension, we could add a conditional at the end, and both the key and value can be any valid expression.