When to Use a List Comprehension in Python

Use for Loops

The most common type of loop is the for loop. You can use a for loop to create a list of elements in three steps:

1. Instantiate an empty list.
2. Loop over an iterable or range of elements.
3. Append each element to the end of the list.

If you want to create a list containing the first ten perfect squares, then you can complete these steps in three lines of code:

>>> squares = []
>>> for number in range(10):
...     squares.append(number * number)
...
>>> squares
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

Here, you instantiate an empty list, squares. Then, you use a for loop to iterate over range(10). Finally, you multiply each number by itself and append the result to the end of the list.

Work With map Objects

For an alternative approach that’s based in functional programming, you can use map(). You pass in a function and an iterable, and map() will create an object. This object contains the result that you’d get from running each iterable element through the supplied function.

As an example, consider a situation in which you need to calculate the price after tax for a list of transactions:

>>> prices = [1.09, 23.56, 57.84, 4.56, 6.78]
>>> TAX_RATE = .08
>>> def get_price_with_tax(price):
...     return price * (1 + TAX_RATE)
...

>>> final_prices = map(get_price_with_tax, prices)
>>> final_prices
<map object at 0x7f34da341f90>

>>> list(final_prices)
[1.1772000000000002, 25.4448, 62.467200000000005, 4.9248, 7.322400000000001]

Here, you have an iterable, prices, and a function, get_price_with_tax(). You pass both of these arguments to map() and store the resulting map object in final_prices. Finally, you convert final_prices into a list using list().

Leverage List Comprehensions

List comprehensions are a third way of making or transforming lists. With this elegant approach, you could rewrite the for loop from the first example in just a single line of code:

>>> squares = [number * number for number in range(10)]
>>> squares
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

Rather than creating an empty list and adding each element to the end, you simply define the list and its contents at the same time by following this format:

new_list = [expression for member in iterable]

Every list comprehension in Python includes three elements:

1. expression is the member itself, a call to a method, or any other valid expression that returns a value. In the example above, the expression number * number is the square of the member value.
2. member is the object or value in the list or iterable. In the example above, the member value is number.
3. iterable is a list, set, sequence, generator, or any other object that can return its elements one at a time. In the example above, the iterable is range(10).

Because the expression requirement is so flexible, a list comprehension in Python works well in many places where you would use map(). You can rewrite the pricing example with its own list comprehension:

>>> prices = [1.09, 23.56, 57.84, 4.56, 6.78]
>>> TAX_RATE = .08
>>> def get_price_with_tax(price):
...     return price * (1 + TAX_RATE)
...

>>> final_prices = [get_price_with_tax(price) for price in prices]
>>> final_prices
[1.1772000000000002, 25.4448, 62.467200000000005, 4.9248, 7.322400000000001]

The only distinction between this implementation and map() is that the list comprehension in Python returns a list, not a map object.

Filter Values From a List

The most common way to add conditional logic to a list comprehension is to add a conditional to the end of the expression.

Earlier, you saw this formula for how to create list comprehensions:

new_list = [expression for member in iterable]

While this formula is accurate, it’s also a bit incomplete. A more complete description of the comprehension formula adds support for optional conditionals.

Here, your conditional statement comes just before the closing bracket:

new_list = [expression for member in iterable if conditional]

Conditionals are important because they allow list comprehensions to filter out unwanted values, which would normally require a call to filter():

>>> sentence = "the rocket came back from mars"
>>> [char for char in sentence if char in "aeiou"]
['e', 'o', 'e', 'a', 'e', 'a', 'o', 'a']

In this code block, the conditional statement filters out any characters in sentence that aren’t vowels.

The conditional can test any valid expression. If you need a more complex filter, then you can even move the conditional logic to a separate function:

>>> sentence = (
...     "The rocket, who was named Ted, came back "
...     "from Mars because he missed his friends."
... )
>>> def is_consonant(letter):
...     vowels = "aeiou"
...     return letter.isalpha() and letter.lower() not in vowels
...

>>> [char for char in sentence if is_consonant(char)]
['T', 'h', 'r', 'c', 'k', 't', 'w', 'h', 'w', 's', 'n', 'm', 'd',
 'T', 'd', 'c', 'm', 'b', 'c', 'k', 'f', 'r', 'm', 'M', 'r', 's', 'b',
 'c', 's', 'h', 'm', 's', 's', 'd', 'h', 's', 'f', 'r', 'n', 'd', 's']

Here, you create a complex filter, is_consonant(), and pass this function as the conditional statement for your list comprehension. Note that you also pass the member value char as an argument to your function.

You can place the conditional at the end of the statement for basic filtering, but what if you want to change a member value instead of filtering it out? In this case, it’s useful to place the conditional near the beginning of the expression. You can do so by taking advantage of the conditional expression:

new_list = [true_expr if conditional else false_expr for member in iterable]

By placing the conditional logic at the beginning of a list comprehension, you can use conditional logic to select from multiple possible output options. For example, if you have a list of prices, then you may want to replace negative prices with 0 and leave the positive values unchanged:

>>> original_prices = [1.25, -9.45, 10.22, 3.78, -5.92, 1.16]
>>> [price if price > 0 else 0 for price in original_prices]
[1.25, 0, 10.22, 3.78, 0, 1.16]

Here, your expression is a conditional expression, price if price > 0 else 0. This tells Python to output the value of price if the number is positive, but to use 0 if the number is negative. If this seems overwhelming, then it may be helpful to view the conditional logic as its own function:

>>> def get_price(price):
...     return price if price > 0 else 0
...

>>> [get_price(price) for price in original_prices]
[1.25, 0, 10.22, 3.78, 0, 1.16]

Now, your conditional expression is contained within get_price(), and you can use it as part of your list comprehension.

Remove Duplicates With Set and Dictionary Comprehensions

While the list comprehension in Python is a common tool, you can also create set and dictionary comprehensions. A set comprehension is almost exactly the same as a list comprehension in Python. The difference is that set comprehensions make sure the output contains no duplicates. You can create a set comprehension by using curly braces instead of brackets:

>>> quote = "life, uh, finds a way"
>>> {char for char in quote if char in "aeiou"}
{'a', 'e', 'u', 'i'}

Your set comprehension outputs all the unique vowels that it found in quote. Unlike lists, sets don’t guarantee that items will be saved in any particular order. This is why the first member of the set is a, even though the first vowel in quote is i.

Dictionary comprehensions are similar, with the additional requirement of defining a key:

>>> {number: number * number for number in range(10)}
{0: 0, 1: 1, 2: 4, 3: 9, 4: 16, 5: 25, 6: 36, 7: 49, 8: 64, 9: 81}

To create the dictionary, you use curly braces ({}) as well as a key-value pair (number: number * number) in your expression.

Assign Values With the Walrus Operator

Python 3.8 introduced the assignment expression, also known as the walrus operator. To understand how you can use it, consider the following example.

Say you need to make ten requests to an API that will return temperature data. You only want to return results that are greater than 100 degrees Fahrenheit. Assume that each request will return different data. In this case, the formula expression for member in iterable if conditional provides no way for the conditional to assign data to a variable that the expression can access.

You need the temperature in both the expression and the conditional so this is a challenge. The walrus operator (:=) solves this problem. It allows you to run an expression while simultaneously assigning the output value to a variable. The following example shows how this is possible, using get_weather_data() to generate fake weather data:

>>> import random
>>> def get_weather_data():
...     return random.randrange(90, 110)
...

>>> [temp for _ in range(20) if (temp := get_weather_data()) >= 100]
[107, 102, 109, 104, 107, 109, 108, 101, 104]

Note that the walrus operator needs to be in the conditional part of your comprehension. You won’t often need to use the assignment expression inside of a list comprehension in Python, but it’s a useful tool to have at your disposal when necessary.

Watch Out for Nested Comprehensions

You can nest comprehensions to create combinations of lists, dictionaries, and sets within a collection. For example, say a climate laboratory is tracking the high temperature in five different cities for the first week of June. The perfect data structure for storing this data could be a Python list nested within a dictionary. You can create the data using nested comprehensions:

>>> cities = ["Austin", "Tacoma", "Topeka", "Sacramento", "Charlotte"