Comprehensions and generators
List, dict and set comprehensions and when they stop being readable; generator expressions and yield; why a generator can read a file larger than memory.
Comprehensions are the most recognisably Python thing in the language, and generators are the idea that makes large data tractable. Both are asked about, and generators are asked about badly enough that a clear answer stands out.
Comprehensions
A comprehension is a loop that produces a collection, written as one expression:
# the loop
squares = []
for n in range(5):
squares.append(n * n)
# the comprehension
squares = [n * n for n in range(5)] # [0, 1, 4, 9, 16]Read it left to right as "n squared, for each n in range 5". Add a filter with if:
evens = [n for n in range(10) if n % 2 == 0] # [0, 2, 4, 6, 8]The same syntax builds the other three collection types:
{n: n * n for n in range(4)} # dict {0: 0, 1: 1, 2: 4, 3: 9}
{c for c in "placement"} # set unique characters
(n * n for n in range(4)) # generator — NOT a tupleThat last line is the one people get wrong. Brackets give a list, braces give a set or dict, and
parentheses give a generator, not a tuple. To build a tuple, call tuple(...) on a generator.
Nesting reads outer-loop-first, which is the opposite of what most people guess:
grid = [[1, 2], [3, 4]]
flat = [x for row in grid for x in row] # [1, 2, 3, 4]Compare with the loop it replaces and the order is obvious:
flat = []
for row in grid: # first
for x in row: # second
flat.append(x)A conditional expression goes before the for; a filter goes after it. Both in one line is
where readability usually dies:
[n if n % 2 == 0 else 0 for n in range(5)] # [0, 0, 2, 0, 4] — transform
[n for n in range(5) if n % 2 == 0] # [0, 2, 4] — filterThe honest limit. A comprehension that needs a comment is worse than the loop it replaced.
Two for clauses and one if is about the ceiling. Interviewers reading your code care more about
whether they can follow it than whether it is one line.
Generators
A comprehension builds the whole result in memory. A generator produces values one at a time, on demand, and remembers where it was.
squares_list = [n * n for n in range(1_000_000)] # ~40 MB
squares_gen = (n * n for n in range(1_000_000)) # a few hundred bytesBoth can be iterated once with a for. Only the list can be indexed, measured with len, or
iterated twice:
gen = (n for n in range(3))
print(sum(gen)) # 3
print(sum(gen)) # 0 — exhausted, nothing leftA generator is single-pass. That is the property to say out loud, because it is the one that causes real bugs.
yield
The other way to make a generator is a function containing yield. Calling it does not run the
body — it returns a generator object. The body advances only when the next value is requested:
def countdown(n):
print("starting")
while n > 0:
yield n
n -= 1
print("done")
gen = countdown(3) # nothing printed yet
print(next(gen)) # "starting" then 3
print(next(gen)) # 2
print(list(gen)) # [1] then "done"yield suspends the function, hands a value back, and resumes from the same line next time — with
all its local variables intact. That is the whole idea.
Where it earns its place
Reading a file that does not fit in memory:
def read_lines(path):
with open(path) as f:
for line in f: # the file object is itself lazy
yield line.rstrip("\n")
long_lines = sum(1 for line in read_lines("big.log") if len(line) > 200)Memory use here is one line, whether the file is 2 KB or 20 GB.
An infinite sequence, which a list cannot represent at all:
def fibonacci():
a, b = 0, 1
while True:
yield a
a, b = b, a + b
from itertools import islice
print(list(islice(fibonacci(), 10)))
# [0, 1, 1, 2, 3, 5, 8, 13, 21, 34]And short-circuiting: any and all stop at the first decisive value, so passing them a generator
avoids computing the rest:
if any(is_prime(n) for n in candidates): # stops at the first prime
...Note the missing brackets. any([...]) would build the entire list first, defeating the point.
Iterator versus iterable
A precise pair of definitions, because this is a standard question:
- An iterable is anything you can loop over — it implements
__iter__. Lists, strings, dicts, files, ranges. - An iterator is the object that does the walking — it implements
__next__and remembers its position.iter(x)gives you one.
xs = [1, 2, 3] # iterable
it = iter(xs) # iterator
print(next(it)) # 1
print(next(it)) # 2
print(next(it)) # 3
print(next(it)) # StopIterationA for loop is exactly this: call iter(), call next() repeatedly, stop on StopIteration.
A list is iterable but not an iterator — which is why you can loop over it twice. A generator is both its own iterable and its own iterator, which is why you cannot.
The functions that consume lazily
sum(n * n for n in range(100))
max(students, key=lambda s: s[1])
any(x < 0 for x in xs)
all(x > 0 for x in xs)
sorted(gen) # must consume everything — sorting needs all of it
enumerate(xs) # lazy
zip(a, b) # lazy
map(str.upper, words) # lazy
filter(None, values) # lazy, drops falsy valuesIn Python 3, map, filter and zip all return lazy iterators, not lists. Printing one shows an
object, not values, and this trips people up:
print(map(str, [1, 2])) # <map object at 0x...>
print(list(map(str, [1, 2]))) # ['1', '2']Most of the time a comprehension is clearer than map or filter and costs the same.
What to take into an interview
[...]list,{...}set or dict,(...)generator — not a tuple.- Nested comprehensions read in the same order as the nested loops they replace.
- A generator is lazy and single-pass; a list is eager and reusable.
yieldsuspends a function and resumes it with its locals intact.- Generators make constant-memory processing of arbitrarily large input possible, and infinite sequences representable.
- An iterable can produce an iterator; an iterator holds the position.
foruses both. map,filterandzipare lazy in Python 3.