Python Comprehensions are Awesome!

When programming, it's quite common to want to create a collection of some kind, from another collection, usually with some modification taking place along the way. Python gives an awesome set of tools for dealing with this kind of problem: comprehensions. If you're not using comprehensions regularly in your code, read on, and I'll show you what you're missing out on!

Basic list comprehensions

So, let's create a simple scenario where we might want to iterate over a list in order to create a new list.

numbers = [1, 2, 3, 4, 5]
doubled = []

for number in numbers:
  doubled.append(number * 2)

Simple enough. We grab each number one at a time, and then double that number as we append it to a new list called doubled.

To be clear, this code is totally functional, and the syntax is all fine. However, we can write a pattern like this in a much shorter way, without sacrificing readability. Sounds good right?

Enter list comprehensions:

numbers = [1, 2, 3, 4, 5]
doubled = [number * 2 for number in numbers]

That's really all there is to it. Just take a moment and admire how succinct and readable that is. It's genuinely one of my favourite pieces of Python syntax.

The structure of a list comprehension is as follows:

new_list = [ <value to append> <for loop definition> ]

Since the second half of the comprehension is just the for loop definitions we're all used to writing, the comprehension syntax should come very easily to you. Just like with for loops, we can iterate over any iterable: it doesn't have to be a list.

The next time you find yourself appending to a list while iterating over another collection, stop and see if you can write that structure as a comprehension instead.

Getting fancy with conditionals

So, what happens if we have some conditional logic inside our for loop, and we're only appending some of the values to this new list. That's a fairly common use case, so can we deal with this inside comprhensions? Absolutely!

Let's look at a new example, where we have a list of names, and we want to create a new list which only contains names which end in "n".

names = ["Anne", "Paul", "Martin", "Helen", "Jake", "Alan", "Sarah"]
final_n_names = []

for name in names:
  if name.endswith("n"):
    final_n_names.append(name)

If you're not familiar with the endswith method, it just returns True or False depending on whether a given string matches the suffix provided as an argument.

Again, we have some perfectly functional code here, but it's also a bit more verbose than we really need. Instead, we can once again use a comprehension.

names = ["Anne", "Paul", "Martin", "Helen", "Jake", "Alan", "Sarah"]
final_n_names = [name for name in names if name.endswith("n")]

As you can see, the condition is just tacked on the end of our for loop definition. We can therefore update our comprehension syntax to look like this:

new_list = [ <value to append> <for loop definition> <filtering conditions> ]

Multiple conditions

Eagle eyed readers may have noticed that I wrote filtering conditions in the plural. This is because we can actually specify as many conditions as we want. We can either chain together a series of if statements, or we can use boolean operators like and & or to make complex conditions.

The sky's your limit! Just don't make it too complicated, because things can get hard to read. If your comprehension is getting too long, you're probably going to want to use a for loop instead.

Readability is king!

Conditions first

We can also write a conditional statement at the start of the comprehension, but this means something very different. Remember that the first thing in our list comprehension is the value to append to our new list, so if we use a condition here, what happens?

We end up adding a boolean value!

names = ["Anne", "Paul", "Martin", "Helen", "Jake", "Alan", "Sarah"]
final_n_names = [name.endswith("n") for name in names]

# [False, False, True, True, False, True, False]

Probably not quite what we wanted.

Multiple loops in a single comprehension

Just as we're not limited to a single if statement inside a comprehension, we can also use multiple loop definitions. Be careful with this, because it can quickly get hard to parse, but it's still a useful thing to know about.

For example, we can specify a pair of loops inside a comprehension to generate all of the possible roll combinations for two six sided dice:

roll_combos = [(d1, d2) for d1 in range(1, 7) for d2 in range(1, 7)]

In this comprehension, the items in our final list are actually all tuples. As you can see, we append (d1, d2) to the new list for each iteration of the loops.

The output for the line above is as follows:

[(1, 1), (1, 2), (1, 3), (1, 4), (1, 5), (1, 6), (2, 1), (2, 2), (2, 3), (2, 4), (2, 5), (2,6), (3, 1), (3, 2), (3, 3), (3, 4), (3, 5), (3, 6), (4, 1), (4, 2), (4, 3), (4, 4), (4, 5), (4, 6), (5, 1), (5, 2), (5, 3), (5, 4), (5, 5), (5, 6), (6, 1), (6, 2), (6, 3), (6, 4), (6, 5), (6, 6)]

Not bad for a single line of code.

Different types of comprehension

Set comprehensions

Lists comprehenesions are awesome, but they're not the only kind of comprehesion we have at our disposal. For example, we can do a set comprehension instead. The syntax is absolutely identical: we just wrap the comprehension in curly braces {} instead of square brackets [].

numbers = [1, 2, 3, 4, 5]
doubled = {number * 2 for number in numbers}

The result of a set comprehension is, as you might imagine, a set.

Dictionary comprehensions

Another less common type of comprehension is a dictionary comprehension.

Dictionary comprehensions look like set comprehensions, except we provide a key and a value as the "append value". You're likely to see a dictionary comprehension in combination with zip, but unless you need to perform some modifications to the keys and values, you're better of using the dict constructor instead.

Of course, we don't need to use zip, and we can create a dictionary comprehension by iterating over a simple list.

Here is an example where we take a list of names and use them as keys. We then assign the length of those names the value for each key.

names = ["Anne", "Paul", "Martin", "Helen"]
name_dict = {name.lower(): len(name) for name in names}

# {'anne': 4, 'paul': 4, 'martin': 6, 'helen': 5}

Again, dictionary comprehensions are a lot more niche than something like a list comprehension, but they're still a valuable tool to know about. They might just be the exact thing you need for a given problem.

Wrapping up

Hopefully you're now excited to use comprehensions in your own code! They really are one of my favourite things in Python, and they do so much to cut down on boilerplate code. Very often they end up more readable too!

Just don't get carried away. Sometimes we need a plain old for loop, especially when we're performing complicated logic. Keep is simple. Keep it readable.

Until next time!