List Comprehensions
Here’s a situation you may have run into: you need to build a list (either from an existing list or some data source) but need to apply some transformations to each item in the list. Not a problem for us, we can easily use a for loop!
# Lets say I have a list of colors
colors = ['red', 'blue', 'green', 'yellow', 'purple']
# And a function that returns True if that color is a primary color
def is_primary(color):
primary_colors = ['red', 'blue', 'yellow']
return color in primary_colors
# Now I want to filter out primary colors from the list
# I also want to transform any non-primary color to be all uppercase
uppercase_non_primary_colors = []
for color in colors:
if not is_primary(color):
uppercase_non_primary_colors.append(color.upper())
print(uppercase_non_primary_colors) # Output should be ['GREEN', 'PURPLE']
OK this does work as we expect it to, but Python gives us a really nice tool for situations exactly like this where were applying transformations and filters to our calculated list.
Instead, we could get uppercase_non_primary_colors like this:
uppercase_non_primary_colors = [color.upper() for color in colors if not is_primary(color)]
Theres a lot going on here in one line, so lets break it down a bit. A list comprehension has the following components:
[TRANSFORMATIONS for VARS in SOURCE if CONDITION]
TRANSFORMATIONS: This is where you can modify the final result for each item of the resulting list (in our main example, this is uppercasing the color withcolor.upper())VARS: This is just us declaring the variable name we want to use in the transformations and later conditions (in the example, we need to declare that the variable name iscolorso that we can uppercase it in the transformation)SOURCE: This is the “source data” that the comprehension is building the final result from. This doesn’t just have to be only a List either, it can be any qualifying Iterable (and later we’ll see we can also do comprehensions on Dicts as well!)CONDITION: This is the filter we can apply on each item in the final list. It’s worth noting that this is technically optional and if you dont want to filter anything, just leave out theifpart completely! (Example:[color for color in colors]is a valid comprehension)
A common example you’ll see online when people talk about list comprehensions is calculating square roots of numbers. Lets see what that might look like:
# Here is our list of numbers
nums = [1, 2, 3, 4, 5]
# Calculate the square root of each number
squared_nums = [x ** 2 for x in nums]
print(squared_nums) # Output: [1, 4, 9, 16, 25]
Dict Comprehensions
Comprehensions also work with Dicts as well! Instead of using [], we use {} just like for normal Dicts:
# Lets say we have a dict of prices
prices = { "apple": 0.5, "banana": 0.3, "orange": 0.8 }
# We want to apply a discount to all prices over 0.5
discounted_prices = { item: (price - 0.5) for item, price in prices.items() if price > 0.5 }
print(discounted_prices) # Output: {'orange': 0.3}
The syntax is mainly the same with two key differences: in the TRANSFORMATIONS we now have an additional : to separate what will be come the key and value in the final dict.
Additionally in the VARS part of the comprehension, we had to define two variables, one for the source dict keys and one for the values. (In the example, item is the key and price is the value)
When to use comprehensions
List and Dict Comprehensions are a powerful tool in our Python toolbelt. However, deeply nesting and combining multiple comprehensions together can quickly create code that is difficult to read.
In general, I recommend keeping comprehensions simple: only use them when you have either a single transformation or filter, or one of both. Beyond that, I’ve found that its trying to quickly read code that is combining multiple comprehensions together a headache. So for the sake of code cleanliness, keep it simple.
It’s also worth considering reaching for a library like NumPy for when those complex data transformation scendarios come up.