Comprehensions and generator expressions in Python
By John Lekberg on July 16, 2020.
This week's post is about using comprehensions and generator expressions in Python. You will learn:
 How to use comprehensions to create lists, sets, and dictionaries.
 How to use generator expressions to keep memory usage low.
What are comprehensions?
Comprehensions are a way to create lists, sets, and dictionaries by transforming and filtering other iterables. E.g.
[ (x ** 2) for x in [1, 1, 2, 3, 5] ]
[1, 1, 4, 9, 25]
{ word.casefold(): word for word in ["Hello", "THERE", "geneRAL"] }
{'hello': 'Hello', 'there': 'THERE', 'general': 'geneRAL'}
real_commands = { "sit", "stay", "heel", "wait" } issued_commands = { "SIT", "Sit", "Bark", "jump", "heel", "BARK" } { command for command in issued_commands if command.lower() not in real_commands }
{'BARK', 'Bark', 'jump'}
Here's how I would write the above code without using comprehensions.

Instead of
[ (x ** 2) for x in [1, 1, 2, 3, 5] ]
[1, 1, 4, 9, 25]
I would write
a = [] for x in [1, 1, 2, 3, 5]: a.append(x ** 2) a
[1, 1, 4, 9, 25]

Instead of
{ word.casefold(): word for word in ["Hello", "THERE", "geneRAL"] }
{'hello': 'Hello', 'there': 'THERE', 'general': 'geneRAL'}
I would write
d = {} for word in ["Hello", "THERE", "geneRAL"]: d[word.casefold()] = word d
{'hello': 'Hello', 'there': 'THERE', 'general': 'geneRAL'}

Instead of
real_commands = { "sit", "stay", "heel", "wait" } issued_commands = { "SIT", "Sit", "Bark", "jump", "heel", "BARK" } { command for command in issued_commands if command.lower() not in real_commands }
{'BARK', 'Bark', 'jump'}
I would write
real_commands = { "sit", "stay", "heel", "wait" } issued_commands = { "SIT", "Sit", "Bark", "jump", "heel", "BARK" } s = set() for command in issued_commands: if command.lower() not in real_commands: s.add(command) s
{'BARK', 'Bark', 'jump'}
There are three types of comprehensions:

List comprehensions (PEP 202) create lists and look like
[ exp forin ... ]
E.g.
[ x ** 2 for x in range(4, 5) ]
[16, 9, 4, 1, 0, 1, 4, 9, 16]

Set comprehensions create sets and look like
{ exp forin ... }
E.g.
{ x ** 2 for x in range(4, 5) }
{0, 1, 4, 9, 16}

Dictionary comprehensions (PEP 274) create dicts and look like
{ exp: exp forin ... }
E.g.
{ x: x ** 2 for x in range(4, 5) }
{ 4: 16, 3: 9, 2: 4, 1: 1, 0: 0, 1: 1, 2: 4, 3: 9, 4: 16 }
Using multiple forloops and ifstatements in a comprehension
Comprehensions can have multiple forloops and ifstatements. E.g. Here is a list comprehension that generates some Pythagorean triples:
domain = range(1, 100) [ (a, b, c) for a in domain for b in domain for c in domain if a < b if b < c if a ** 2 + b ** 2 == c ** 2 ]
[(3, 4, 5),
(5, 12, 13),
(6, 8, 10),
...
(57, 76, 95),
(60, 63, 87),
(65, 72, 97)]
Without using a list comprehension, I would write this as:
domain = range(1, 100) triples = [] for a in domain: for b in domain: for c in domain: if a < b: if b < c: if a ** 2 + b ** 2 == c ** 2: triples.append((a, b, c)) triples
[(3, 4, 5),
(5, 12, 13),
(6, 8, 10),
...
(57, 76, 95),
(60, 63, 87),
(65, 72, 97)]
The forloops and ifstatements can be mixed together. Doing this can lead to speedups:

Compare 26.6 seconds
import cProfile domain = range(1, 500) cProfile.run(""" [ (a, b, c) for a in domain for b in domain for c in domain if a < b if b < c if a ** 2 + b ** 2 == c ** 2 ] """)
4 function calls in 26.598 seconds ...

... to 23.0 seconds.
import cProfile domain = range(1, 500) cProfile.run(""" [ (a, b, c) for a in domain for b in domain if a < b for c in domain if b < c if a ** 2 + b ** 2 == c ** 2 ] """)
4 function calls in 23.055 seconds ...
What are generator expressions?
Generator expressions (PEP 289) are a way to create generators (a kind of iterator) by transforming and filtering other iterables. Think of generator expressions as "iterator comprehensions". E.g. Here's a generator expression of the first few square numbers:
( x ** 2 for x in range(10) )
<generator object <genexpr> at 0x10511cc80>
Generator expressions have a similar syntax to comprehensions. They look like
( exp forin ... )
To access the data in the generator, I need to consume the iterator using, e.g., list:
list(( x ** 2 for x in range(10) ))
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
NOTE: When a generator expression is the only argument to a function, the parentheses do not need to be written:
list(( x ** 2 for x in range(10) ))
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
list(x ** 2 for x in range(10))
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
What's the use of generator expressions?

I can generate "comprehensions" for collections like tuple and frozenset:
tuple(x ** 2 for x in range(10))
(0, 1, 4, 9, 16, 25, 36, 49, 64, 81)
frozenset(x ** 2 for x in range(10))
frozenset({0, 1, 4, 9, 16, 25, 36, 49, 64, 81})

Generator expressions don't load all the elements into memory, making it easier to work on large amounts of data. E.g. Here's a function that reports object memory usage:
import bisect import sys def report_memory(exp): x = eval(exp) size_bytes = sys.getsizeof(x) coefficients = [1024 ** n for n in [1, 2]] units = ["B", "KiB", "MiB"] idx = bisect.bisect(coefficients, size_bytes) coef = [1, *coefficients][idx] unit = units[idx] print(f"{exp!r} takes {size_bytes//coef} {unit}")
Compare 819 mebibytes:
domain = range(100_000_000) report_memory('[ x ** 2 for x in domain ]')
'[ x ** 2 for x in domain ]' takes 819 MiB
... to 112 bytes:
report_memory('( x ** 2 for x in domain )')
'( x ** 2 for x in domain )' takes 112 B
Because generator expressions can keep memory usage lower than comprehensions, I like to use them with these functions:
In conclusion...
In this week's post, you learned about comprehensions and generator expressions, which are a concise way to create lists (and other containers) by transforming and subsetting iterables. The comprehension "notation" is similar to setbuilder notation in mathematics.
My challenge to you:
I have a 3 byte passcode with this SHA256 checksum:
6d6125cc4538aaec9dbef490ab1091a6cb4af5348f96a5cb0bfeeeda6edfebbe
Use a comprehension or generator expression to figure out what 3 byte passcode produces this checksum.
You can generate the checksum using hashlib.sha256:
from hashlib import sha256 sha256(b"Hello World").hexdigest()
'a591a6d40bf420404a011733cfb7b190d62c65bf0bcda32b57b277d9ad9f146e'
If you enjoyed this week's post, share it with your friends and stay tuned for next week's post. See you then!
(If you spot any errors or typos on this post, contact me via my contact page.)