List comprehensions

Python has three types of comprehensions: list, dictionary, and set comprehensions.

List comprehensions take a generator expression, an iteration expression, and an optional predicate:

[generator iteration predicate]

The iteration expression must contain an iterable! The generator expression is applied to each item in the iteration expresion. If there is a predicate, then the predicate is evaluated for each item in the iteration expression, and the generation expression is applied to the items for which the predicate evaulates to true. Here is an example of a simple squared list comprehension:

x = [i**2 for i in range(5)]
print(x)
# Output:
# [0, 1, 4, 9, 16]

This is equivalent to:

x = []
for i in range(5):
  x.append(i**2)
print(x)
# Output:
# [0, 1, 4, 9, 16]

Obviously, the first form is much preferable. Here is an example with a predicate:

x = [i for i in range(5) if i % 2 == 0]
print(x)
# Output:
# [0, 2, 4]

The non-comprehension version:

x = []
for i in range(5):
  if i % 2 == 0:
    x.append(i)
print(x)
# Output:
# [0, 2, 4]

You can nest list comprehensions to iterate arbitrarily deep in a structure. Here is an example of a two-deep list comprehension:

x = [1, 2, 3]
y = [4, 5, 6]
l = [i*j for j in x for i in y]
print(l)
# Output:
# [4, 5, 6, 8, 10, 12, 12, 15, 18]

This is equivalent to:

x = [1, 2, 3]
y = [4, 5, 6]
l = []
for j in x:
  for i in y:
    l.append(i*j)
print(l)
# Output:
# [4, 5, 6, 8, 10, 12, 12, 15, 18]

You can make nested comprehensions easier to understand with good indentation:

x = [1, 2, 3]
y = [4, 5, 6]
l = [i*j
     for j in x
     for i in y]
print(l)
# Output:
# [4, 5, 6, 8, 10, 12, 12, 15, 18]

Note that this closely resembles the nested for loop, with the body put at the top.

Dictionary comprehensions

Dictionary comprehensions are very similar to list comprehensions:

x = [1, 2, 3, 4, 5]
d = {i: i**2 for i in x}
print(d)
# Output:
# {1: 1, 2: 4, 3: 9, 4: 16, 5: 25}
This is equivalent to:

x = [1, 2, 3, 4, 5]
d = {}
for i in x:
  d[i] = i**2
print(d)

Set comprehensions

Set comprehensions follow the same rules as dictionary and list comprehensions:

a = [1, 1, 1, 2, 3, 4, 4, 5, 6]
{x for x in a}
print(x)
# Output:
# {1, 2, 3, 4, 5, 6}

Mix and match

You can combine forms to create arbitrarily complex expressions. Here is an example of a list comprehension and dictionary comprehension that reverses the keys and values of a 1-deep dictionary:

x = [1, 2, 3, 4, 5]
d = {i: i**2 for i in x}
z = {v: k for k, v in d.items()}
print(d)
print(z)
# Output:
# {1: 1, 2: 4, 3: 9, 4: 16, 5: 25}
# {1: 1, 4: 2, 9: 3, 16: 4, 25: 5}

Building a list of keys with more than one value:

d = {'foo': [1, 2, 3], 'bar': [1], 'baz': [1, 4, 0]}
l = [k for k in d if len(x[k]) > 1]
print(l)
# Output:
# ['foo', 'baz']

Building a set of unique keys from all objects in the dictionary:

d = {
    'foo': {
        'attr1': 'foobar',
        'attr2': 'barfoo'
    },
    'bar': {
        'attr1': 'foobaz',
        'attr3': 'bazbuf'
    },
    'baz': {
        'attr2': 'bufbuf',
        'attr4': 'bufbaz'
    }
}

s = {key for item in d.values() for key in item.keys()}
print(s)
# Output:
# {'attr2', 'attr1', 'attr4', 'attr3'}

References