Comprehensions are one of those mildly advanced features everyone should learn, but shouldn’t be used everywhere. They are more efficient, but they can hurt readability if overly complex. Comprehensions come in three main flavors:
- List:
[var for var in iterable] - Dictionary:
{key: value for key, value in iterable} - Generator:
(var for var in iterable)
I’ll focus on list comprehensions for examples since they are one of the most common, but dictionary and generator comprehensions are extremely useful for manipulating data.
Nested Comprehensions
One thing I wanted to touch on is comprehensions can be nested together. This roughly translates to a double for loop, and it’s easier to think of the iteration as following that pattern just inlined.
These two chunks of code, one written with a.append and a double for loop, and the other written with a nested list comprehension, are identical in function.
a = []
for i in range(10):
for j in range(10):
a.append((i, i * j))
As a nested comprehension:
a = [(i, i * j) for i in range(10) for j in range(10)]
Comprehensions can be nested well past two. I once took a bet with someone that I could write a very heavily looping structure as a nested comprehension. I was correct, it could be done, but the resulting code was far less readable.
When can you use comprehensions?
Comprehensions can be used when you’re doing minimal logic inside of a loop to build a list, dictionary, or generator. For example, say we have a piece of user input we need to transform with another function before returning it. You might write:
results = []
for piece in user_input:
results.append(process_input(piece))
This can be quickly transformed into a comprehension by moving things around a little bit:
results = [process_input(piece) for piece in user_input]
Let’s use some WORD placeholders instead of code to make this a bit simpler. The original code looked like this:
results = []
for VARIABLE in ITERABLE:
results.append(EXPRESSION)
You can translate this as:
results = [EXPRESSION for VARIABLE in ITERABLE]
What else can comprehensions do?
Comprehensions can also contain a small amount of logic. Let’s say you wanted to process a list of strings, but you only want ones which are not empty. You might write a loop like this:
results = []
for piece in input:
if piece:
results.append(piece)
This can also be translated into a comprehension by using the if <condition> term, which will skip (continue) the current variable if it is not met.
results = [piece for piece in input if piece]
And using more abstract terms as above:
results = []
for VARIABLE in ITERABLE:
if CONDITION:
results.append(EXPRESSION)
You can translate this as:
results = [EXPRESSION for VARIABLE in ITERABLE if CONDITION]
When shouldn’t you use a comprehension?
You shouldn’t use a comprehension when there are multiple expressions inside the loop, when you’re holding or maintaining extra variables or state, or when it would make the logic less clear about the actual intent. The last piece is highly subjective. You can feel the comprehension is very clear at the time it’s written but it could quickly become less clear, especially if you’ve nested them.
My rule of thumb is if I find myself adding comments to different “steps” of the comprehension it is too complex.