Python Comprehensions

1. We can use a for loop to do a kind of mapping. Suppose we want to double every number in a list of numbers.

     def double_list(ns):
         doubles = []
         for n in ns:
             doubles.append(2 * n)
         return doubles
   
   >>> double_list([1,2,3])
   [2, 4, 6]
   

2. We can combine for with if to filter out all negative numbers from a list of numbers

   def drop_negatives(ns):
       nonnegs = []
       for n in ns:
           if n >= 0:
               nonnegs.append(n)
       return nonnegs
   
   >>> drop_negatives([-2, -1, 0, 1, 2])
   [0, 1, 2]
   

3. We can nest for loops

   def cross_product(xs, ys):
       tuples = []
       for x in xs:
           for y in ys:
               tuples.append((x, y))
       return tuples
   
   >>> cross_product([1,2,3], [4,5])
   [(1, 4), (1, 5), (2, 4), (2, 5), (3, 4), (3, 5)]
   

4. And we can combine nested looping with a filter to get all pairs (2-tuples) where the first element is positive and the second negative.

   def posnegs(xs, ys):
       pns = []
       for x in xs:
           for y in ys:
               if x > 0 and y < 0:
                   pns.append((x,y))
       return pns
   
   >>> posnegs([1, -2, 3, -4], [5, -6, 7, -8])
   [(1, -6), (1, -8), (3, -6), (3, -8)]
   

A comprehension is an abbreviated way to write such combinations of mapping and filtering with if and for. It turns the code from the previous section inside out and obviates the need for a temporary variable.

We will first look at comprehensions over lists in some detail and then see how to make comprehensions for sets and dictionaries.

1. We can write the list of doubles as

   def double_list(ns):
       return [2 * n for n in ns]
   

   The list comprehension has several advantages:

   1. It is more compact.
   2. It expresses directly the list we want.
   3. There is no need to first set a variable and then update it; this takes place behind the scenes.
   4. The loop variable doubles is not left 'lying around' after the for loop. Although there are cases where we might make use of the value of the loop variable after completion of the for loop, we generally do not and so, strictly speaking, should del it.

2. Similarly we can drop the negatives from a list as

   def drop_negatives(ns):
       return [n for n in ns if n >= 0]
   
   >>> drop_negatives([-2, -1, 0, 1, 2])
   [0, 1, 2]
   

3. The cross product becomes

   def cross_product(xs, ys):
       return [(x,y) for x in xs for y in ys]
   
   >>> cross_product([1,2,3], [4,5])
   [(1, 4), (1, 5), (2, 4), (2, 5), (3, 4), (3, 5)]
   

4. And the nested loop with filter becomes

   def posnegs(xs, ys):
       return [(x,y) for x in xs for y in ys if x > 0 and y < 0]
   
   >>> posnegs([1, -2, 3, -4], [5, -6, 7, -8])
   [(1, -6), (1, -8), (3, -6), (3, -8)]
   

   Note how the deeply nested if within a for within a for gets flattened into a succinct, readable form.

5. To aid readability further the clauses can be placed under one another

   def posnegs(xs, ys):
       return [(x, y)
               for x in xs
               for y in ys
               if x > 0 and y < 0]
   

6. Although all of the above examples are on lists of numbers, comprehensions work on lists of all data types, e.g. strings. If we are given a list of strings and asked to make a list of the lengths of those strings, we might use a for loop as follows:

   strings = ['a', 'bc', 'def', 'ghijklmn']
   
   lengths = []
   for string in strings:
       lengths.append(len(string))
   

   Now lengths contains the list

   [1, 2, 3, 8]
   

   The list comprehension compacts the for loop into a single line

   lengths = [len(string) for string in strings]
   

   This can be read as "the list of len(string) with string taking, in turn, the value of each element of strings".

7. Here is an example of filtering a list of numbers and taking the square root of only the positive numbers.

   import math
   
   math.sqrt(4)
   
   ns = [1, -2, 3, -4, 5, -6]
   roots = [math.sqrt(n) for n in ns if n > 0]
   

   In this example we first import the math library and show how to use the sqrt function. Then we return a list of the square roots of only the positive elements of the list ns.

   Using a for loop instead of a comprehension, we might write:

   roots = []
   for n in ns:
       if n > 0:
           roots.append(math.sqrt(n))
   

8. This example flattens a list of lists of numbers.

   elements = [x for l in [[1,2,3], [4,5,6], [7,8,9]] for x in l]
   

9. More examples can be found in the documentation.

In the same way as we can write list comprehensions, so too can we write set comprehensions. We simply write braces {} instead of square brackets []

{x for x in [1,2,3,2,1]}
{c for c in 'abacadabra'}

Note how the latter example iterates over the characters in a string.

Filtering, nesting and multiple iteration are also possible.

>>> {letter for letter in 'abstemiously' if letter not in 'aeiou'}
{'b', 'l', 'm', 's', 't', 'y'}

In the same way as we write set comprehensions, so too can we write dictionary comprehensions, for example to map strings to their lengths.

{s:len(s) for s in ['a', 'bc', 'def', 'ghijklm']}

An equivalent for loop is:

d = {}
for s in ['a', 'bc', 'def']:
    d[s] = len(s)

We can also map across a dictionary

d = {"a": 3, "b": 5, "c": 7, "d": 9}
[k for k in d if d[k] > 6]

Since we can write comprehensions for lists, sets and dictionaries, one might wonder whether we can write tuple comprehensions. With a moment's thought we realise tuples are immutable so we cannot build a tuple as follows:

# NOT ALLOWED! tuples are immutable
t = ()
for x in [1, 2, 3]:
    t.append(x)

And if we try to build a tuple using a comprehension, we get a strange looking result:

>>> (x for x in [1,2,3])
<generator object <genexpr> at 0x7f28be212620>

This result is actually a generator, an object that we can iterate over. Generators are dealt with in the chapter on Control Flow