list comprehensions

The Python list comprehensions are a very easy way to apply a function or filter to a list of items. List comprehensions can be very useful if used correctly but very unreadable if you're not careful.

Let's dive right into a few examples. The basic premise of a list comprehension looks like this:

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

We can easily expand this with a filter:

>>> uneven_squares = [x ** 2 for x in range(10) if x % 2]
>>> uneven_squares
[1, 9, 25, 49, 81]

The syntax is pretty close to regular Python for loops, but the if statement and automatic storing of results makes it quite useful for some cases. The regular Python equivalent is not much longer, however:

>>> uneven_squares = []
>>> for x in range(10):
... if x % 2:
... uneven_squares.append(x ** 2)

>>> uneven_squares
[1, 9, 25, 49, 81]

Care must be taken though; because of the special list comprehension structure, some types of operations are not as obvious as you might expect. This time, we are looking for random numbers greater than 0.5:

>>> import random
>>> [random.random() for _ in range(10) if random.random() >= 0.5]
[0.5211948104577864, 0.650010512129705, 0.021427316545174158]

See that last number? It's actually less than 0.5. This happens because the first and the last random calls are actually separate calls and return different results.

One way to counter this is by creating the list separate from the filter:

>>> import random
>>> numbers = [random.random() for _ in range(10)]
>>> [x for x in numbers if x >= 0.5]
[0.715510247827078, 0.8426277505519564, 0.5071133900377911]

That obviously works, but it's not all that pretty. So what other options are there? Well, there are a few but the readability is a bit questionable, so these are not the solutions that I would recommend. It's good to see them at least once, however.

Here is a list comprehension in a list comprehension:

>>> import random
>>> [x for x in [random.random() for _ in range(10)] if x >= 0.5]

And here's one that quickly becomes an incomprehensible list comprehension:

>>> import random
>>> [x for _ in range(10) for x in [random.random()] if x >= 0.5]

Caution is needed with these options as the double list comprehension actually works like a nested for loop would, so it quickly generates a lot of results. To elaborate on this regard:

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

This effectively does the following:

>>> results = []
>>> for x in range(3):
... for y in range(3, 5):
... results.append((x, y))
...
>>> results
[(0, 3), (0, 4), (1, 3), (1, 4), (2, 3), (2, 4)]

These can be useful for some cases, but I would recommend that you limit their usage, as they have a tendency to quickly become unreadable. I would strongly advise against using list comprehensions within list comprehensions for the sake of readability. It's still important to understand what is happening, so let's look at one more example. The following list comprehension swaps the column and row counts, so a 3 x 4 matrix becomes 4 x 3:

>>> matrix = [
... [1, 2, 3, 4],
... [5, 6, 7, 8],
... [9, 10, 11, 12],
... ]

>>> reshaped_matrix = [
... [
... [y for x in matrix for y in x][i * len(matrix) + j]
... for j in range(len(matrix))
... ]
... for i in range(len(matrix[0]))
... ]

>>> import pprint
>>> pprint.pprint(reshaped_matrix, width=40)
[[1, 2, 3],
 [4, 5, 6],
 [7, 8, 9],
 [10, 11, 12]]

Even with the extra indentation, the list comprehension just isn't all that readable. With four nested loops, that is expectedly so, of course. There are rare cases where nested list comprehensions might be justified, but generally I won't recommend their usage.