Python Comprehensions: What They Are and Some Cool Tricks

Written by Dan Sackett on October 23, 2014

One of my favorite things about Python is that you can take for loops and if statements and combine them into a compact comprehension to do the same logic.

What am I talking about?

Python comprehensions come in a few types, but each of them allows us to build primitive Python data types in one concise expression. Think back to code where you create an empty list item and then loop through another list of items, use an if condition, and build a new list programatically. Two days ago I showed you how we can avoid this with dictionaries but today I wanted to go deeper into the subject.

In Python, we have three types of comprehensions. We have list comprehensions, generator comprehensions, and dictionary comprehensions. As you can see from the naming, these of course relate to the appropriate Python types. For instance, a list comprehension will return a list. In the case of a list and generator comprehension, we will see the following format:

element for variable(s) in list if condition

Let's walk through this.

As you can see, we have a for loop and an if statement and they can combine in one line to give us a new item. Let's see how this works with lists.

List Comprehensions

List comprehensions are used a lot in Python. They take the form from above and are surrounded by list syntax:

my_list = [element for variable(s) in list if condition]

Before showing this off though, let's show a naive solution to creating a new list.

new_list = []
my_list = [1, 2, 3, 4, 5]

for number in my_list:
    if number < 4:

print new_list
# [1, 2, 3]

In our example, we have a list of number 1-5 and we want to make a new list of numbers less than 4. It's a simple exercise but one that shows us that writing this takes some code. It requires two levels of nesting and requires the append() function to do so. Let's see how we can do this with a list comprehension:

my_list = [1, 2, 3, 4, 5]
new_list = [number for number in my_list if number < 4]

print new_list
# [1, 2, 3]

As you can see, we ditched the nesting and combined our loop and our if statement into a list. This is a comprehension. We're saying "Make me a list of numbers where the number is less than 4 based on my_list". If this looks like magic to you, break it down into three parts like we did above:

Using this concise syntax is one of the features that makes Python very strong and efficient. Let's see another example with integers:

nums = [1, 2, 3, 4, 5]
print [x * x for x in nums]
# [1, 4, 9, 16, 25]

In this example, you can see that our element can be manipulated before it is returned. In this case, we square the number by multiplying it by itself. This results in a list of squares. How about creating a list of squares for only even numbers:

nums = [1, 2, 3, 4, 5]
print [x * x for x in nums if x % 2 == 0]
# [4, 16]

We use the modulo notation to check if the number is even. Since only two of these numbers match that case, our list returned is of those two items only.

While these are all basic computations, take note that this syntax can be extended further to do more complex things. Let's take a matrix for example. We can flatten it with the following script:

my_list = [
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]

print [item for row in my_list for item in row]
# [1, 2, 3, 4, 5, 6, 7, 8, 9]

Notice in this example that we have two for loops. Comprehensions allow us to use as many loops as we want in the syntax. Just be sure that you don't make it unreadable in the end. Good code is concise, not because it is on one line. We can stack if statements as well:

my_list = [
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]

print [item for row in my_list for item in row if item > 2 if item < 7]
# [3, 4, 5, 6]

Hopefully you can start to see some of the power that comprehensions have. With most things, these can get much more complex. Let's look at something a bit more complex:

print [(x,y) for x in range(5) if x % 2 == 0 for y in range(5) if y % 2 == 1]
# [(0, 1), (0, 3), (2, 1), (2, 3), (4, 1), (4, 3)]

In this example we have multiple conditions, and multiple loops running. In the end, we can break it down into two pieces. The first piece loops through the numbers 0-4 and if it's an even number then we go into the second loop. In the second loop, we iterate through numbers 0-4 but if it's an odd number then we accept it. In the end, we create a list of tuples where even numbers are paired with odd numbers in numbers 0-4.

List comprehensions can be a little odd at first, but I recommend experimenting with them. Take a for loop and reduce it into a comprehension. In most cases, you'll find that it's possible and will clean your code up quite a bit while still allowing readability.

Generator Expressions

Generator comprehensions (better know as generator expressions) are pretty much exactly like list comprehensions. If you're new to generators, check out my post on them and why they're so efficient.

my_gen = (element for variable(s) in list if condition)

Looking at that, the only difference is that we use a parenthesis instead of the normal list syntax. This is good news because all of the above examples can be translated the same.

What's the difference then?

Well this goes back to why lists and generators are different:

So naturally, this comes down to efficiency. If you want to write a comprehension for a huge data set, think about what it will cost to store that whole list in memory. It might be smarter to use a generator in this case.

One of the best places to use these are when you are working with functions inline. For instance:

print sum(x**2 for x in range(10))
# 285

In this example, we want to use the sum() function to add up all the squares from 0-9. We could write a function to handle this, but the syntax is much more efficient. In essence, it will compute each value one at a time. Since we don't care about the generator after the sum is computed, this is the perfect solution because generators are exhausted after one iteration. We can move on and forget about it.

This is where the distinction lies. If you don't care about the resulting data set, use a generator expression. If you need the resulting data set, use a list comprehension. If you follow that guideline, you'll be writing efficient code.

Dictionary Comprehensions

Python has a lot of little known features and one of them in my mind is the dictionary comprehension. Like above, they can simplify creating and manipulating a dictionary. For instance, let's take a dictionary and reverse it so they keys become the values. In naive code, we would do the following:

my_dict = {'name': 'Dan', 'gender': 'Male', 'age': 25}
new_dict = {}

for key, value in my_dict.items():
    new_dict[value] = key

print new_dict
# {25: 'age', 'Dan': 'name', 'Male': 'gender'}

We can do this much more efficiently with a comprehension:

my_dict = {'name': 'Dan', 'gender': 'Male', 'age': 25}
print {value: key for key, value in my_dict.items()}
# {25: 'age', 'Dan': 'name', 'Male': 'gender'}

Clean, concise, and efficient. Notice the syntax is very similar to our other examples except for two things:

As long as we follow that, we can build list comprehension examples into dictionary comprehensions.


Comprehensions are one of the features that attracts me to Python. As long as they're used in a smart manner they can reduce your code and still give you a readable result. They're very efficient, especially in the case of a generator expression, and will allow you to think about creating iterable sets in new ways.

Give them a try and see what you can reduce into a comprehension.


comments powered by Disqus