Comprehensions in Python—Write Shorter For Loops

In Python, you can turn your for loops into one-liners by using comprehensions.

This is a comprehensive guide to comprehensions in Python.

Python supports 4 types of comprehension:

1. List comprehensions
2. Dictionary comprehensions
3. Set comprehensions
4. Generator comprehensions

By the way, there are no tuple comprehensions in Python. But you can find a way to mimic tuple comprehensions to get a very similar outcome.

You will learn how each comprehension type works in Python and what are the benefits/drawbacks of using comprehensions in the first place.

1. List Comprehensions

The list comprehension goes with the syntax:

output_list = [expression for var in input_list if condition]

Where the if condition part is optional.

Example

Let’s create a list of numbers by leaving out all the negative numbers. First, you can use a for loop to solve this task:

numbers = [4, -2, 7, -4, 19]
new_nums = []
for num in numbers:
    if num > 0:
        new_nums.append(num)
print(new_nums)

Output:

[4, 7, 19]

But you can make this for loop shorter by using a list comprehension:

new_nums = [num for num in numbers if num > 0]
print(new_nums)

Output:

[4, 7, 19]

2. Dictionary Comprehensions

Python has a shorthand for looping through dictionaries known as dictionary comprehension.

The syntax for dictionary comprehension is:

{ key:value for (key,value) in dict.items() if condition }

Example

Let’s square all the number values of a dictionary using dictionary comprehension:

data = {'a': 1, 'b': 2, 'c': 3, 'd': 4, 'e': 5}
squared = {k:v*v for (k,v) in data.items()}
print(squared)

Output:

{'a': 1, 'b': 4, 'c': 9, 'd': 16, 'e': 25}

3. Set Comprehensions

Set comprehension is like a list comprehension for sets.

The syntax of a set comprehension is:

{ expression for var in input if condition }

Example

Let’s create a list of numbers and pick all the even numbers from it to a set:

numbers = [13, 21, 14, 24, 53, 62] 
  
filtered_nums = set() 
  
for num in numbers: 
    if num % 2 == 0: 
        filtered_nums.add(num) 
  
print(filtered_nums)

Output:

{24, 62, 14}

But using a set comprehension you can do this in one line:

filtered_nums = {num for num in numbers if num % 2 == 0}
print(filtered_nums)

Output:

{24, 62, 14}

4. Generator Comprehensions

Generators support comprehension. Generator comprehension is better known as generator expression.

Similar to the other comprehensions, the generator expression offers a shorthand syntax for a for loop.

The syntax of the generator expression is:

( expression for var in input if condition )

Example

Let’s square all the even numbers of a list and drop all the odds.

First, you can use a for loop to get the job done:

def square_even(numbers):
    for number in numbers:
        if number % 2 == 0:
            yield(number * number)
            
numbers = [1, 2, 3, 4, 5, 6]
squared_numbers = square_even(numbers)

for number in squared_numbers:
    print(number)

Output:

4
16
36

With a generator expression, you can forget the square_even() altogether and do the same with one line of code:

squared_numbers = (num * num for num in numbers if num % 2 == 0)
for number in squared_numbers: 
    print(number)

Output:

4
16
36

When Should You Not Use Comprehensions in Python?

Now you’ve learned the four comprehension types in Python.

It might be tempting to replace all your projects’ for loops with comprehensions. Before doing that, however, let’s discuss when you should not use comprehensions.

In short, do not use comprehensions if it sacrifices code quality!

A good example is when dealing with nested for loops. If you write a nested for loop as a comprehension, you can save lines of code, but the code quality may become worse.

For example, a matrix is often represented as a list of lists in Python. If you want to flatten a matrix, you can use a list comprehension:

matrix = [
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9],
]
 
flat_matrix = [num for row in matrix for num in row]
print(flat_matrix)

Output:

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

The code is concise, but it might not be that intuitive to understand what it does.

If you use a nested for-loop instead, the code becomes more understandable:

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

for row in matrix:
    for num in row:
        flat_matrix.append(num)

print(flat_matrix)

Output:

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

Comprehensions vs For Loops: Performance

Finally, let’s take a look at the performance of a comprehension

For example, a list comprehension is around 10% quicker than a for loop when squaring 1M integers.

Performing similar benchmarking for set and dictionary comprehensions gives somewhat similar results: A comprehension performs slightly better.

(For generators, comparing generator expression and a for-loop + yield is meaningless: Both return a generator object almost instantly. The time is not spent on computing the values as generators do not store values.)

To take home: Don’t use comprehensions for performance. The performance benefit is insignificant compared to the harm caused by lengthy and verbose comprehensions. Write comprehensions when it cleans up your code without losing its quality.

Conclusion

Turn your for loops into one-liners by using comprehensions.

Python supports four comprehensions for the common data structures:

• List Comprehension
• Dictionary Comprehension
• Set Comprehension
• Generator Comprehension

Using comprehension wisely can improve your code quality. But do not blindly use comprehensions. Sometimes replacing a nested for loop with comprehension can reduce the code understandability.

Comprehensions can outperform a for loop. However, do not use comprehensions for performance. Instead, favor the code quality.

Thanks for reading. Happy coding!