Python Comprehension Tutorial

Python Comprehension Tutorial

This Python tutorial will focus on comprehensions and how to use them. The concept of comprehension in Python is that it is an easier and more readable way to create a list or other sequence. Creating a list is the most common application of comprehension, but it can also be used with a dictionary, tuple, set, or even a string. It is also possible to add logic to a comprehension with an if statement or an if else statement. We’ll inspect common for loop structures, and how they can be rewritten using a comprehension in Python.


Simple List Comprehension

A good way to learn comprehensions in Python is by seeing code that uses a for loop versus code that uses a comprehension to get the desired result. We can start with a list comprehension and use a for loop code snippet to begin.

for loop

Here we have a simple list of numbers stored in the numbers variable. We are then using a for loop to loop over each item, and adding it to the my_list[] variable on each iteration. This code kind of reads like, I want a ‘number’ for each item in ‘numbers’.

[1, 2, 3, 4, 5, 6, 7]

list comprehension

The list comprehension version does away with the for loop, and simply uses the [number for number in numbers] syntax to achieve the same result. It reads like, I want a number for each number in numbers. Printing out the my_list variable shows that we have successfully copied each item from numbers to my_list using a comprehension in Python. Cool!

[1, 2, 3, 4, 5, 6, 7]

Taking An Action On Each Item

When using a for loop, it is common to apply some type of action to each item for each iteration. In this example, we are going to calculate the square of each number in the list. You can read it like, I want ‘number * number’ for each ‘number’ in ‘numbers’. Let’s see what this looks like in a for loop.

for loop

To do this in a for loop, we create an empty list then loop through all the numbers contained in it. On each iteration, we append the square of the number. When the for loop stops running, we simply print out the list to see the result.

[1, 4, 9, 16, 25, 36, 49]

list comprehension

The list comprehension version of this loop is very similar to the first example we looked at. The only difference this time around, is that we are saying we want the square for each item. So this reads like, I want ‘number * number’ for each ‘number’ in ‘numbers’ and you can see how well that translates to the comprehension of [number * number for number in numbers]. Printing out the resulting list gives us exactly what we wanted.

[1, 4, 9, 16, 25, 36, 49]

lambda version

This can also be accomplished with a lambda function, but it is not as easy to read or comprehend. I think you’ll agree.

[1, 4, 9, 16, 25, 36, 49]

Comprehension If

When using the for loop, we can use an if condition within the loop to determine if we want to take an action or not.

for loop

For this example, we are now going to square the number on each iteration *only* if the number is even. We can check if a number is even by using the modulo operator with the syntax of if number % 2 == 0.

[4, 16, 36]

list comprehension

The same result can be achieved in the list comprehension version. We can use an if statement in the comprehension to take an action only if a condition is met.

[4, 16, 36]

Comprehension If Else

During a loop it is also common to use an if/else construct to do one action if one condition, and a different action for a different condition.

for loop

In this for loop, we are going to use the if/else to append the number if it is even, and append the square of the number if it is odd.

[1, 2, 9, 4, 25, 6, 49]

list comprehension

In the list comprehension, an if/else can also be used as a one-liner.

[1, 2, 9, 4, 25, 6, 49]

So we are starting to see how we are making efficient use of iteration in Python by simply using the comprehensions syntax.


Comprehension Nested

Nested for loops are used in Python for various problems. You can nest using a comprehension as well.

for loop

First, we look at creating letter and number pairs. For example, we want something like a1, a2, a3, b1, b2, b3, etc… To do this with a for loop, it looks like this.

[('a', 1), ('a', 2), ('a', 3), ('b', 1), ('b', 2), ('b', 3), ('c', 1), ('c', 2), ('c', 3)]

tuple comprehension

Now we translate the nested for to a comprehension on one line with the same result.

[('a', 1), ('a', 2), ('a', 3), ('b', 1), ('b', 2), ('b', 3), ('c', 1), ('c', 2), ('c', 3)]

Comprehensions Dictionary

In addition to creating lists, we can create Python dictionaries using the comprehensions syntax.

for loop

We learned about the zip() function in our Python tips and tricks tutorial. In this for loop, we’ll use that zip() function to associate values from one list with another list. This results in a dictionary of key/value pairs populated using the index value of each list.

{'Red': 'Pepper', 'Green': 'Onion', 'Orange': 'Squash'}

dict comprehension

The same effect can be created using a dictionary comprehension in Python.

{'Red': 'Pepper', 'Green': 'Onion', 'Orange': 'Squash'}

Dict Comprehension With If

Using a conditional statement also works in dictionary comprehensions.

for loop

{'Red': 'Pepper', 'Green': 'Onion'}

dictionary comprehension with if

Here is the dictionary comprehension version which uses an if statement.

{'Red': 'Pepper', 'Green': 'Onion'}

Comprehension Set

A set in Python is a sequence with no duplicate values. Here we will take a list of numbers with several duplicate values, and create a set both using a for loop and a set comprehension.

for loop

{1, 2, 3, 4, 5}

set comprehension

{1, 2, 3, 4, 5}

Learn More About Python Comprehensions

Python Comprehension Tutorial Summary

Comprehensions are another example of the elegant syntax you can use to create effective code with Python. By using comprehensions, many times you can reduce code that previously required two, three, four, or more lines of code with for loops and nested for loops, right down to a one-liner. Once you are used to their syntax, comprehensions are also very readable and usually a better option than using .map() and .filter() functions.