Python Primer — List comprehensions

Hello World,

The previous posts in this Python Primer series dealt with general concepts that are pretty much part of any programming language. We started with data types, then took a look at statements and in the last post functions and objects were covered.

In this post we’ll tackle a topic that is more “pythonic” in nature — List comprehensions.

Concept

List comprehensions are all about making the code shorter and easier to read. When working with lists, they essentially allow for a block of code that would take several lines to be condensed to a one-liner.

So, when creating a list, we can generally do it in a single line. Let’s see a simple example of this.

We have a list with numbers and we want to make another list based on that. Although this example is not practical, it illustrates what list comprehensions are. To see the difference and why comprehensions are useful, first block code will show how we would do this “the old fashioned way”.

initial_list = [2, 7, 27, 28, 44, 49, 63, 64, 84, 100]
new_list = []
for nr in initial_list:
    new_list.append(nr)
print(new_list)

>> [2, 7, 27, 28, 44, 49, 63, 64, 84, 100]

In the above code, we went through each item of the initial_list in the for loop and append each one to the new_list.

Using list comprehensions, we can do this in more elegant, one line:

initial_list = [2, 7, 27, 28, 44, 49, 63, 64, 84, 100]
new_list = [nr for nr in initial_list]
print(new_list)

>> [2, 7, 27, 28, 44, 49, 63, 64, 84, 100]

Let’s break everything down:

1. The new_list variable is created, the square brackets indicating it’s a list. In the list comprehension version, there is no need to declare an empty list first
2. We iterate over the initial list in the for loop. The loop is also inside the list comprehension
3. The current element in the iteration is added in the new_list. In the classic version, the append method must be used. With list comprehensions it’s enough to simply use the variable name, since list comprehensions automatically creates the list

Hopefully the concept of list comprehension is clearer now, both how it works and why it’s preferable to do it all in a single line. But it wouldn’t be that useful if we could only copy elements from one list to another…

Changing variable value

Let’s take the same example list as before, but instead populate the new list with the squares of the original list.

initial_list = [2, 7, 27, 28, 44, 49, 63, 64, 84, 100]
new_list = []
for nr in initial_list:
    new_list.append(nr**2)
print(new_list)

>> [4, 49, 729, 784, 1936, 2401, 3969, 4096, 7056, 10000]

With list comprehension we do it in a single line.

initial_list = [2, 7, 27, 28, 44, 49, 63, 64, 84, 100]
new_list = [nr**2 for nr in initial_list]
print(new_list)

>> [4, 49, 729, 784, 1936, 2401, 3969, 4096, 7056, 10000]

This example shows we can change the element inside a list comprehension, just like we could inside a regular for loop.

Adding an if statement

Let’s presume we want to have all the squares of numbers in the original list, but only if they are even numbers. The conditional if is needed for this. Fortunately, it can be used inside a list comprehension.

initial_list = [2, 7, 27, 28, 44, 49, 63, 64, 84, 100]
new_list = []
for nr in initial_list:
    if nr % 2 == 0:
        new_list.append(nr**2)
print(new_list)

>> [4, 784, 1936, 4096, 7056, 10000]

With list comprehension, the example looks like this:

initial_list = [2, 7, 27, 28, 44, 49, 63, 64, 84, 100]
new_list = [nr**2 for nr in initial_list if nr % 2 == 0]
print(new_list)

>> [4, 784, 1936, 4096, 7056, 10000]

Breakdown time again:

1. The new_list variable is created, the square brackets indicating it’s a list. In the list comprehension version, there is no need to declare it as an empty list first
2. We iterate over the initial list in the for loop. The loop is also inside the list comprehension
3. If the number is even, we append it to the list. In the comprehension version, the if statement is after the for statement
4. The current element in the iteration is added in the new_list. In the classic version, the append method must be used. With list comprehensions version, it’s enough to simply use the variable name, since list comprehensions automatically creates the list.

Using multiple variables

List comprehension can also be used with multiple for loops. To illustrate this, let’s say we want to generate a list of tuples, each tuple containing a 2D coordinate set, from 0 to 4.

coord_list = []
for x in range(4):
    for y in range(4):
        coord_list.append((x,y))
print(coord_list)

>> [(0, 0), (0, 1), (0, 2), (0, 3), (1, 0), (1, 1), (1, 2), (1, 3), (2, 0), (2, 1), (2, 2), (2, 3), (3, 0), (3, 1), (3, 2), (3, 3)]

The list comprehension version of this example looks like this:

coord_list = [(x,y) for x in range(4) for y in range(4)]
print(coord_list)
>> [(0, 0), (0, 1), (0, 2), (0, 3), (1, 0), (1, 1), (1, 2), (1, 3), (2, 0), (2, 1), (2, 2), (2, 3), (3, 0), (3, 1), (3, 2), (3, 3)]

Let’s break down this final example as well:

1. The coord_list variable is created, the square brackets indicating it’s a list. In the list comprehension version, there is no need to declare it as an empty list first
2. We iterate over range(4) and keep each iteration in the x variable, which accounts for the first coordinate value
3. We iterate over range(4) and keep each iteration in the y variable, which accounts for the second coordinate value. The for loops order is preserved in the list comprehension
4. We append a tuple of the current iteration of x and y (x, y) to the list. Inside the list comprehension we write the tuple as is, without the need for the append method

Conclusion

List comprehensions are very useful when writing Python code and a must for any Python developer. There is a caveat however, as some developers tend to over use them and end up with a one line that is too complex to read through, which is exactly against the purpose list comprehensions exist in the first place.

So use them well and use them often at your own judgement but don’t sacrifice readability for the cool factor.

The previous post seemed too long so this won’t be the last post in the Python Primer series after all… The next one will cover lambda functions and where to use them.

Your opinions, feedback or (constructive) criticism are welcomed in discussions below or at @mariussafta

Join our Facebook and Meetup groups and keep an eye out for discussions and future meetups.