Trey Hunner

programming, teaching, open source

Python List Comprehensions: Explained Visually

| Comments

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.

Update: I held a 1 hour video chat about list comprehensions which extends the material in this article. If you want more after reading this post, check out the recording.

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:

1
2
>>> doubled_odds = map(lambda n: n * 2, filter(lambda n: n % 2 == 1, numbers))
>>> doubled_odds = [n * 2 for n in numbers if n % 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:

1
2
3
4
new_things = []
for ITEM in old_things:
    if condition_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 " + ITEM for ITEM in old_things if condition_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:

1
2
3
4
5
6
numbers = [1, 2, 3, 4, 5]

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

To this:

1
2
3
numbers = [1, 2, 3, 4, 5]

doubled_odds = [n * 2 for n in numbers if n % 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 * 2 for n in numbers if 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 * 2 for 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 = [n for row in matrix for 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 = [n for n in row for 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: key for 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 * 2 for n in numbers if n % 2 == 1]

After

1
2
3
4
5
doubled_odds = [
    n * 2
    for n in numbers
    if n % 2 == 1
]

Nested loops in list comprehension

Before

1
flattened = [n for row in matrix for n in row]

After

1
2
3
4
5
flattened = [
    n
    for row in matrix
    for n in row
]

Dictionary comprehension

Before

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

After

1
2
3
4
flipped = {
    value: key
    for key, value in original.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 " + ITEM for ITEM in old_things if 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 article was based on my Intro to Python class. If you’re interested in chatting about my Python training services, drop me a line.

Learn more through weekly Python chats 🎁

Like my teaching style? Want to learn more? Sign up for attend my Weekly Python Chat events so I can answer your questions about Python, programming, and life in general.

Comments