Generators and Iterators
Objectives
By the end of this chapter, you should be able to:
- Explain what a generator is and how to create one in Python
- Explain what the
yieldkeyword andnext()function do - Define and create iterators using the
iter()function - Access values and indexes using
enumerate
💡 Why this matters: Computing an entire dataset upfront wastes memory and time if you’ll only ever need the next value or two. Generators let you produce results lazily, on demand, which matters a lot once you’re working with large or unbounded sequences.
Generators
A generator function looks like a normal function, but instead of return it uses the yield keyword. Calling a generator function doesn’t run its code right away: it hands you back a special generator object that only runs the code when you actually ask for the next value.
Here’s the difference made concrete. First, a normal function with return:
def simple_return():
print("start")
return 1
print("this line never runs")
result = simple_return()
print(result)
# start
# 1
return ends the function immediately, so the line after it never executes. Now the generator version, using yield instead:
def simple_gen():
print("start")
yield 1
print("middle")
yield 2
print("end")
g = simple_gen()
print("generator created, but nothing above has printed yet")
print(next(g))
print(next(g))
# generator created, but nothing above has printed yet
# start
# 1
# middle
# 2
Notice "generator created..." prints before "start" does. Calling simple_gen() doesn’t run any of the function’s code, it just creates the generator object. Each call to next() runs the function until it hits a yield, hands back that value, and pauses right there, remembering exactly where it left off. The next next() call resumes from that paused point instead of starting the function over. That pause-and-resume behavior is what makes yield different from return, which always ends the function for good.
A more typical generator loops and yields multiple values:
def gen_squares(n):
for num in range(n):
yield num ** 2
for x in gen_squares(5):
print(x)
# 0
# 1
# 4
# 9
# 16
Because each value is only computed when it’s asked for, generators are a great fit for large, or even infinite, sequences: you never have to hold the whole thing in memory at once.
Here’s a Fibonacci sequence generator, using tuple unpacking to update both values at once:
def fib_with_generator(n):
a, b = 1, 1
for _ in range(n):
yield a
a, b = b, a + b
for num in fib_with_generator(8):
print(num)
# 1
# 1
# 2
# 3
# 5
# 8
# 13
# 21
A Generator Only Runs Once
A generator object gets used up as you pull values out of it. Once you’ve iterated all the way through, it’s empty, and iterating it again gives you nothing:
def gen_squares(n):
for num in range(n):
yield num ** 2
squares = gen_squares(3)
print(list(squares)) # [0, 1, 4]
print(list(squares)) # [] (already exhausted, nothing left to give)
If you need to loop over the same values twice, call the generator function again to get a fresh generator object. A single generator object can only be consumed once.
The next() Function
Given a generator object, next() retrieves its next value:
def use_next():
for x in range(10):
yield x
gen = use_next()
print(next(gen)) # 0
print(next(gen)) # 1
print(next(gen)) # 2
Call next() past the last value and you’ll get a StopIteration error: the generator has nothing left to give. A for loop handles this automatically, catching StopIteration for you so your program doesn’t crash:
for val in use_next():
print(val)
# 0
# 1
# 2
# ...
# 9
You can also write a generator as a generator expression: like a list comprehension, but with parentheses instead of square brackets:
def use_next():
return (x for x in range(10))
Iterators
Calling iter() on something turns it into an iterator you can step through with next():
word = "hello"
word_iter = iter(word)
print(next(word_iter)) # 'h'
print(next(word_iter)) # 'e'
print(next(word_iter)) # 'l'
print(next(word_iter)) # 'l'
print(next(word_iter)) # 'o'
print(next(word_iter)) # raises StopIteration
An iterator is the more general concept: anything with a __next__ method that produces a next value, one call at a time. Every generator is an iterator (Python builds that __next__ method for you automatically when you write yield), but not every iterator is a generator. A list or string isn’t itself an iterator, which is exactly why you had to call iter() on it above to get one.
enumerate
Sometimes you want both an item and its index while looping. enumerate gives you both, as an (index, value) pair on each iteration:
items = ["first", "second", "third"]
for idx, value in enumerate(items):
print(f"index is {idx} and value is {value}")
# index is 0 and value is first
# index is 1 and value is second
# index is 2 and value is third
all and any
Two built-ins for checking truthiness across an iterable.
all() returns True only if every element is truthy:
print(all([0])) # False
print(all([0, 1])) # False
print(all([0, "", [1]])) # False
print(all([1, "a", [1]])) # True
any() returns True if at least one element is truthy:
print(any([0])) # False
print(any([0, 1])) # True
print(any([0, "", [1]])) # True
For more advanced iteration patterns, the standard library’s itertools module is worth exploring once you’re comfortable here.
Try It
- Write a generator function that yields the first
npowers of 2, and loop over it with aforloop. - Call
next()directly on a generator until it raisesStopIteration, and observe the error. - Use
enumerateto print each character of a string alongside its index.
Recap
- A generator function uses
yieldto produce values lazily, one at a time, instead of computing everything upfront. Unlikereturn,yieldpauses the function and remembers where it left off for the nextnext()call. - A generator is exhausted after you’ve pulled every value out of it; iterating it again gives you nothing.
next()retrieves the next value from a generator or iterator; aforloop handles the eventualStopIterationfor you.- An iterator is the general concept: anything with
__next__. Every generator is an iterator, but not every iterator is a generator. enumerategives you both the index and the value while looping;all/anycheck truthiness across an iterable.
Next lesson: decorators, functions that wrap other functions.