Trey Hunner: Python List Comprehensions: Explained Visually

Sometimes a programming design pattern becomes common enough to warrant its own special syntax. Python’s list comprehensions are a prime example of such a syntactic sugar.

List comprehensions in Python are great, but mastering them can be tricky because they don’t solve a new problem: they just provide a new syntax to solve an existing problem.

Let’s learn what list comprehensions are and how to identify when to use them.

What are list comprehensions?

List comprehensions are a tool for transforming one list (any iterable actually) into another list. During this transformation, elements can be conditionally included in the new list and each element can be transformed as needed.

If you’re familiar with functional programming, you can think of list comprehensions as syntactic sugar for a filter followed by a map:

12

>>> doubled_odds=map(lambdan:n*2,filter(lambdan:n%2==1,numbers))>>> doubled_odds=[n*2forninnumbersifn%2==1]

If you’re not familiar with functional programming, don’t worry: I’ll explain using for loops.

From loops to comprehensions

Every list comprehension can be rewritten as a for loop but not every for loop can be rewritten as a list comprehension.

The key to understanding when to use list comprehensions is to practice identifying problems that smell like list comprehensions.

If you can rewrite your code to look just like this for loop, you can also rewrite it as a list comprehension:

1234

new_things=[]forITEMinold_things:ifcondition_based_on(ITEM):new_things.append("something with "+ITEM)

You can rewrite the above for loop as a list comprehension like this:

1

new_things=["something with "+ITEMforITEMinold_thingsifcondition_based_on(ITEM)]

List Comprehensions: The Animated Movie™

That’s great, but how did we do that?

We copy-pasted our way from a for loop to a list comprehension.

Here’s the order we copy-paste in:

1. Copy the variable assignment for our new empty list (line 3)
2. Copy the expression that we’ve been append-ing into this new list (line 6)
3. Copy the for loop line, excluding the final : (line 4)
4. Copy the if statement line, also without the : (line 5)

We’ve now copied our way from this:

123456

numbers=[1,2,3,4,5]doubled_odds=[]forninnumbers:ifn%2==1:doubled_odds.append(n*2)

To this:

123

numbers=[1,2,3,4,5]doubled_odds=[n*2forninnumbersifn%2==1]

List Comprehensions: Now in Color

Let’s use colors to highlight what’s going on.

doubled_odds = []for n in numbers:
    if n % 2 == 1:
        doubled_odds.append(n * 2)

doubled_odds = [n * 2for n in numbersif n % 2 == 1]

We copy-paste from a for loop into a list comprehension by:

1. Copying the variable assignment for our new empty list
2. Copying the expression that we’ve been append-ing into this new list
3. Copying the for loop line, excluding the final :
4. Copying the if statement line, also without the :

Unconditional Comprehensions

But what about comprehensions that don’t have a conditional clause (that if SOMETHING part at the end)? These loop-and-append for loops are even simpler than the loop-and-conditionally-append ones we’ve already covered.

A for loop that doesn’t have an if statement:

doubled_numbers = []for n in numbers:
    doubled_numbers.append(n * 2)

That same code written as a comprehension:

doubled_numbers = [n * 2for n in numbers]

Here’s the transformation animated:

We can copy-paste our way from a simple loop-and-append for loop by:

1. Copying the variable assignment for our new empty list (line 3)
2. Copying the expression that we’ve been append-ing into this new list (line 5)
3. Copying the for loop line, excluding the final : (line 4)

Nested Loops

What about list comprehensions with nested looping?… 😦

Here’s a for loop that flattens a matrix (a list of lists):

flattened = []for row in matrix:
    for n in row:
        flattened.append(n)

Here’s a list comprehension that does the same thing:

flattened = [nfor row in matrixfor n in row]

Nested loops in list comprehensions do not read like English prose.

Note: My brain wants to write this list comprehension as:

flattened = [nfor n in rowfor row in matrix]

But that’s not right! I’ve mistakenly flipped the for loops here. The correct version is the one above.

When working with nested loops in list comprehensions remember that the for clauses remain in the same order as in our original for loops.

Other Comprehensions

This same principle applies to set comprehensions and dictionary comprehensions.

Code that creates a set of all the first letters in a sequence of words:

first_letters = set()for w in words:
    first_letters.add(w[0])

That same code written as a set comprehension:

first_letters = {w[0]for w in words}

Code that makes a new dictionary by swapping the keys and values of the original one:

flipped = {}for key, value in original.items():
    flipped[value] = key

That same code written as a dictionary comprehension:

flipped = {value: keyfor key, value in original.items()}

Readability Counts

Did you find the above list comprehensions hard to read? I often find longer list comprehensions very difficult to read when they’re written on one line.

Remember that Python allows line breaks between brackets and braces.

List comprehension

Before

1

doubled_odds=[n*2forninnumbersifn%2==1]

After

12345

doubled_odds=[n*2forninnumbersifn%2==1]

Nested loops in list comprehension

Before

1

flattened=[nforninrowforrowinmatrix]

After

12345

flattened=[nforrowinmatrixforninrow]

Dictionary comprehension

Before

1

flipped={value:keyforkey,valueinoriginal.items()}

After

1234

flipped={value:keyforkey,valueinoriginal.items()}

Note that we are not adding line breaks arbitrarily: we’re breaking between each of the lines of code we copy-pasted to make these comprehension. Our line breaks occur where color changes occur in the colorized versions.

Learn with me

I did a class on list comprehensions with PyLadies Remote recently.

If you’d like to watch me walk through an explanation of any of the above topics, check out the video:

1. list comprehensions
2. generator expressions
3. set comprehensions
4. dictionary comprehensions

Summary

When struggling to write a comprehension, don’t panic. Start with a for loop first and copy-paste your way into a comprehension.

Any for loop that looks like this:

new_things = []for ITEM in old_things:
    if condition_based_on(ITEM):
        new_things.append("something with " + ITEM)

Can be rewritten into a list comprehension like this:

new_things = ["something with " + ITEMfor ITEM in old_thingsif condition_based_on(ITEM)]

If you can nudge a for loop until it looks like the ones above, you can rewrite it as a list comprehension.

This blog post was based on my Intro to Python class. If you’re interested in chatting about my Python training services, drop me a line.