How to Use Comprehensions in Python

Introduction

In Python, comprehension refers to a concise syntax for creating new sequences from existing iterables.

They can be applied to any iterable, which includes lists, sets, and dictionaries, among others. But they will also work if we create our own iterables.

The basic syntax is:

expression for element in iterable

Optionally, we can add a filter condition like this:

expression for element in iterable if condition

This syntax has minimal variations depending on the object we want to generate with the comprehension (it generally varies in whether we wrap it with [] or {}).

But let’s see it better with some examples 👇

List Comprehensions

List comprehensions allow creating new lists by applying an expression to each element of a sequence. The basic syntax is:

[expression for element in iterable if condition]

For example, we can apply it without a condition:

# Create a list of squares of numbers from 0 to 9
squares = [x**2 for x in range(10)]

print(squares)  # Output: [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

Or in this example with a condition:

# Create a list of squares of the numbers 0 to 9 that are even
even_squares = [x**2 for x in range(10) if x % 2 == 0]

print(even_squares)  # Output: [0, 4, 16, 36, 64]

# Create a 3x3 identity matrix
identity_matrix = [[1 if i == j else 0 for j in range(3)] for i in range(3)]

print(identity_matrix)
# Output: [[1, 0, 0], [0, 1, 0], [0, 0, 1]]

Set Comprehensions

Set comprehensions are similar to list comprehensions, but they generate a set. The syntax is:

{expression for element in iterable if condition}

For example:

# Create a set of odd numbers from 0 to 9
odds = {x for x in range(10) if x % 2 != 0}

print(odds)  # Output: {1, 3, 5, 7, 9}

Dictionary Comprehensions

We can also use comprehensions that allow us to build dictionaries concisely. The syntax is:

{key: value for element in iterable if condition}

For example, without a condition:

# Create a dictionary with numbers and their squares for numbers from 0 to 9
squares_dict = {x: x**2 for x in range(10)}

print(squares_dict)
# Output: {0: 0, 1: 1, 2: 4, 3: 9, 4: 16, 5: 25, 6: 36, 7: 49, 8: 64, 9: 81}

Or with a condition:

# Create a dictionary of squares only for even numbers from 0 to 9
even_squares_dict = {x: x**2 for x in range(10) if x % 2 == 0}

print(even_squares_dict)
# Output: {0: 0, 2: 4, 4: 16, 6: 36, 8: 64}

Generator Comprehensions

Generators return an object that produces elements. The syntax is:

(expression for element in iterable if condition)

For example:

# Create a generator of squares of numbers from 0 to 9
squares_gen = (x**2 for x in range(10))

print(list(squares_gen))  # Output: [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

It may seem similar to the list case, but in this case, what we have is a generator iterable.

This is useful if we work with large volumes of data that we don’t need to generate and load into memory all at once.

That is, in the list example, squares had all the values calculated. However, in the case of squares_gen, each value is calculated only once each time we iterate over the collection.

In the example, they seem similar because print(list()) iterates over all elements and prints them.