Iterators and generators in Python — practice with endless Fibonacci, a prime pipeline and the library booking queue (and the laptop that kept thinking)

For a class exercise I needed “a few” Fibonacci numbers and, just in case, I built a list of the first ten million. The laptop kept thinking, the fan took off and in the end I didn’t even use them: I only needed the ones up to 100. The lesson I took away is the one behind this whole article: don’t compute in advance what you don’t know you will need; ask for values one at a time and stop when you have enough.

In the iterators and generators theory you saw how the for loop works inside, why yield pauses instead of ending and why a generator runs out. Today we put it to work with three medium-level programs: an infinite Fibonacci generator, a pipeline of prime number generators and the booking queue for a library book, built with linked nodes and walked with a generator. All three use data that looks like the real thing: sequences that never end, limits you don’t know in advance and a queue that changes while you use it.

Before you start

It helps to be clear about:

• A for calls iter() and then next() until StopIteration appears.
• A function with yield returns a generator and runs nothing until the first next().
• return inside a generator ends it (and the for walking it simply finishes).
• A generator can be walked only once; a list, as often as you like.

If any of these sounds shaky, go back to the theory before carrying on.

Program 1: infinite Fibonacci generator

The Fibonacci sequence (0, 1, 1, 2, 3, 5, 8…) is the perfect example of a sequence with no end: each number is the sum of the two before it. The program has to give the first 10, the ones up to 100, the even ones up to 500, the first one that reaches 1000 (with its position) and show how the ratio between two consecutive numbers approaches the golden ratio. What makes it tricky: each question needs a different amount of numbers, and we don’t know how many in advance.

The idea

A single infinite generator, fibonacci(), which only keeps two numbers (the previous and the current one) and hands over one on each next(). Everything else is a different way of stopping asking: take(n, gen) cuts by count, up_to(limit, gen) cuts by value (with return), and a for with break cuts on any condition. The generator does not know who is using it or how many will be asked for; it only knows how to give the next one.

The code

def fibonacci():
    """Infinite generator of the Fibonacci sequence: 0, 1, 1, 2, 3, 5..."""
    previous, current = 0, 1
    while True:
        yield previous
        previous, current = current, previous + current   # the right side is computed BEFORE assigning


def up_to(limit, gen):
    """Lets the values of gen through while they do not exceed the limit."""
    for x in gen:
        if x > limit:
            return                # return in a generator = done (StopIteration)
        yield x


def take(n, gen):
    """The first n values of any generator (like itertools.islice)."""
    for _ in range(n):
        yield next(gen)


print('First 10: ', list(take(10, fibonacci())))
print('Up to 100:', list(up_to(100, fibonacci())))
evens = [n for n in up_to(500, fibonacci()) if n % 2 == 0]
print('Evens up to 500:', evens)

for pos, n in enumerate(fibonacci()):          # enumerate also asks one at a time
    if n >= 1000:
        print(f'First >= 1000: {n} (position {pos})')
        break

g = fibonacci()
next(g)                                        # skip the 0 so we never divide by 0
a = next(g)
for i in range(1, 21):
    b = next(g)
    if i in (5, 10, 20):
        print(f'Ratio after {i:2} steps: {b / a:.6f}')
    a = b

The most important line is previous, current = current, previous + current. Python computes the whole right-hand side before assigning anything: it first builds the tuple (current, previous + current) with the old values and then unpacks it. If you wrote it in two lines (previous = current and then current = previous + current), the second would use the already-changed previous and the sequence would come out wrong.

In the drawing, each violet pill is the generator’s state at a next() and the coral label is what it hands over. The generator never keeps more than two numbers.

Both helper functions are generators too:

def up_to(limit, gen):
    """Lets the values of gen through while they do not exceed the limit."""
    for x in gen:
        if x > limit:
            return                # return in a generator = done (StopIteration)
        yield x


def take(n, gen):
    """The first n values of any generator (like itertools.islice)."""
    for _ in range(n):
        yield next(gen)

• up_to walks the generator it receives and, as soon as a value exceeds the limit, it does return. That ends up_to (and the list() walking it finishes), even though fibonacci() could keep giving numbers forever.
• take calls next(gen) n times and hands over each value. It is what itertools.islice(gen, n) does.

Output

First 10:  [0, 1, 1, 2, 3, 5, 8, 13, 21, 34]
Up to 100: [0, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89]
Evens up to 500: [0, 2, 8, 34, 144]
First >= 1000: 1597 (position 17)
Ratio after  5 steps: 1.600000
Ratio after 10 steps: 1.618182
Ratio after 20 steps: 1.618034

Let’s check it by hand. The first ten: 0, 1, 0 + 1 = 1, 1 + 1 = 2, 1 + 2 = 3, 2 + 3 = 5, 3 + 5 = 8, 13, 21, 34. Up to 100 two more arrive (55 and 89); the next one, 144, goes over and up_to returns. The even ones up to 500 are 0, 2, 8, 34 and 144: the next even one would be 610 (in Fibonacci, one number in three is even). The first one reaching 1000 is 1597 = 610 + 987, at position 17 counting from 0. And the ratios 8 / 5 = 1.6, then 1.618182 and 1.618034: the golden ratio (1.6180339…) appears by itself.

next() no.	previous	current	Hands over	Next state
1	0	1	0	(1, 0 + 1) = (1, 1)
2	1	1	1	(1, 1 + 1) = (1, 2)
3	1	2	1	(2, 1 + 2) = (2, 3)
4	2	3	2	(3, 2 + 3) = (3, 5)
5	3	5	3	(5, 3 + 5) = (5, 8)

Questions you might ask

Why does each question create a new fibonacci()? Because a generator runs out. If we reused the same one for “up to 100” and then for “evens up to 500”, the second would start where the first stopped (at 144) and we would skip the beginning.

Why does up_to use return and not break? Inside a generator, return ends it: it is the clean way to say “there is nothing more”. A break would leave the inner for and, since nothing comes after it, the generator would end too; it would work the same, but return makes it clearer what happens.

Why do we skip the 0 before computing the ratios? Because the ratio is b / a and the first value is 0: dividing by 0 would raise ZeroDivisionError. A lone next(g) throws that first value away.

Why doesn’t enumerate(fibonacci()) hang if the generator is infinite? Because enumerate is lazy too: it asks one at a time. The break in the for stops asking at position 17, and nobody computes a single number more.

Try it yourself

• Write sum_up_to(limit), which returns the sum of the Fibonacci numbers not exceeding the limit, using sum() with up_to() and without creating any list.
• Change fibonacci() so it receives the first two values (fibonacci(2, 1) gives the Lucas sequence: 2, 1, 3, 4, 7…).

Program 2: prime number generator

A prime is a number greater than 1 that can only be divided by 1 and itself. The program has to give the first 15, how many there are up to 100, the 50th prime, the ones between 100 and 150 and the first pairs of twin primes (two primes that differ by 2, such as 11 and 13). What makes it tricky: checking whether a number is prime takes work, the sequence never ends and each question cuts it at a different place.

The idea

Split the responsibilities into small pieces that plug into each other like a pipeline: is_prime(n) only decides whether a number is prime; primes() is the infinite generator that keeps trying candidates; primes_up_to(limit), n_primes(n) and twin_primes(gen) are generators that consume primes() and transform or cut what comes out. Each number travels the whole pipeline before the next one is computed.

The code

def is_prime(n):
    """True if n is prime. Only tries odd divisors up to the square root."""
    if n < 2:
        return False
    if n % 2 == 0:
        return n == 2             # 2 is the only even prime
    i = 3
    while i * i <= n:             # if n has a divisor, one of them is <= root of n
        if n % i == 0:
            return False
        i += 2
    return True


def primes():
    """Infinite generator of primes."""
    yield 2
    candidate = 3
    while True:
        if is_prime(candidate):
            yield candidate
        candidate += 2            # even numbers (except 2) need no checking


def primes_up_to(limit):
    for p in primes():            # consumes the infinite generator one by one
        if p > limit:
            return
        yield p


def n_primes(n):
    counter = 0
    for p in primes():
        if counter == n:
            return
        yield p
        counter += 1


def twin_primes(gen):
    """Pairs of primes that differ by 2, taken from a generator of primes."""
    previous = next(gen)
    for p in gen:                 # the for carries on where next() left it
        if p - previous == 2:
            yield (previous, p)
        previous = p


print('First 15:', list(n_primes(15)))
print('Up to 100 there are', sum(1 for _ in primes_up_to(100)), 'primes')

g = primes()
for _ in range(49):
    next(g)
print('The 50th prime:', next(g))

print('Between 100 and 150:', [p for p in primes_up_to(150) if p >= 100])

twins = twin_primes(primes())
print('First 5 twin prime pairs:', [next(twins) for _ in range(5)])

The is_prime function has two shortcuts that save a lot of work:

def is_prime(n):
    """True if n is prime. Only tries odd divisors up to the square root."""
    if n < 2:
        return False
    if n % 2 == 0:
        return n == 2             # 2 is the only even prime
    i = 3
    while i * i <= n:             # if n has a divisor, one of them is <= root of n
        if n % i == 0:
            return False
        i += 2
    return True

• Even numbers are settled in one line: return n == 2 gives True only for 2 and False for the rest.
• Only odd divisors up to the square root are tried (i * i <= n). If n = a × b, one of the two factors is at most √n: if there is no divisor up to there, there is none after. For 97 it is enough to try 3, 5, 7 and 9 (11 × 11 = 121 is already over): 4 divisions instead of 95.

And twin_primes uses a detail of the protocol that is worth a close look:

def twin_primes(gen):
    """Pairs of primes that differ by 2, taken from a generator of primes."""
    previous = next(gen)
    for p in gen:                 # the for carries on where next() left it
        if p - previous == 2:
            yield (previous, p)
        previous = p

First it calls next(gen) by hand to have 2 as previous. Then the for p in gen carries on where that next() left it: it receives 3, 5, 7… and does not start over, because gen is an iterator (the same object), not a list.

In the drawing, 149 travels the whole pipeline; 151 does not get past primes_up_to(150), which returns and stops everything.

Output

First 15: [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47]
Up to 100 there are 25 primes
The 50th prime: 229
Between 100 and 150: [101, 103, 107, 109, 113, 127, 131, 137, 139, 149]
First 5 twin prime pairs: [(3, 5), (5, 7), (11, 13), (17, 19), (29, 31)]

Checking by hand: up to 100 there are 25 primes (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97). The 50th prime is 229: the loop’s 49 next() calls swallow the first 49 and the next one is the one we print. Between 100 and 150 there are 10 (notice that 119 = 7 × 17 is not there, and neither is 121 = 11 × 11). The first five twin pairs: 3 and 5, 5 and 7 (5 is in two pairs), 11 and 13, 17 and 19, 29 and 31.

twin_primes lap	previous	p received	p - previous == 2?	Hands over
(before the for)	2	—	—	—
1	2	3	No (1)	—
2	3	5	Yes	(3, 5)
3	5	7	Yes	(5, 7)
4	7	11	No (4)	—
5	11	13	Yes	(11, 13)

Questions you might ask

Why sum(1 for _ in primes_up_to(100)) and not len(...)? Because a generator has no len(): it does not know how many items it will give until it gives them. The generator expression adds a 1 for each prime without creating any list. len(list(primes_up_to(100))) also works, but it builds a list just to count it.

Why does n_primes compare counter == n before the yield? So as not to compute a single extra prime. If we checked after the yield, with n equal to 15 the generator would compute prime number 16 before realising it was already done.

What would happen if primes_up_to used break instead of return? The same: nothing comes after the for, so the function would end. But if tomorrow someone adds code after the loop, with break it would run and with return it wouldn’t. return says exactly what we mean: “this generator is finished”.

Why [next(twins) for _ in range(5)] and not list(twins)? Because there are infinitely many twin primes (or so it is suspected: it is an open problem in mathematics). list(twins) would never end.

Try it yourself

• Write factors(n), a generator that hands over the prime factors of n using primes() (for 360: 2, 2, 2, 3, 3, 5).
• Write gaps(gen), which hands over the distance between each prime and the next, and find the largest gap below 1000.

Program 3: iterator for the booking queue (linked list)

When a library book is on loan, members can book it and join a queue: whoever booked first gets it when it comes back. The program keeps that queue as a linked list: each member sits in a node that points to the next. It has to be able to book (add at the end), serve (remove the first one), say how many there are, send each member a notice with their place and answer in, list() and str(). What makes it tricky: the structure is not a Python list, so the for does not know how to walk it… until we give it an __iter__.

The idea

Each node has member and next. The queue keeps two references: _first (where members are served) and _last (where bookings are added, so we don’t have to walk the whole queue to reach the end). To walk it, __iter__ is a generator that starts at _first, hands over the member of each node and jumps to next until it finds None. Since each call to __iter__ creates a new generator, everything else (notices, __str__, in, list()) comes for free and can be repeated.

The code

class BookingQueue:
    """Linked list with the members waiting for a book, in order of arrival."""

    class _Node:
        def __init__(self, member):
            self.member = member
            self.next = None        # the link to the node behind

    def __init__(self, title):
        self.title = title
        self._first = None
        self._last = None           # keep the last one to append without walking
        self._n = 0

    def reserve(self, member):
        node = self._Node(member)
        if self._first is None:     # empty queue: the new one is first and last
            self._first = node
        else:
            self._last.next = node
        self._last = node
        self._n += 1

    def serve(self):
        """The book is back: the first in the queue leaves."""
        if self._first is None:
            raise IndexError('no bookings')
        node = self._first
        self._first = node.next
        if self._first is None:     # now empty: there is no last either
            self._last = None
        self._n -= 1
        return node.member

    def __len__(self):
        return self._n

    def __iter__(self):             # generator: walks the nodes without touching them
        node = self._first
        while node is not None:
            yield node.member
            node = node.next

    def notices(self):
        for place, member in enumerate(self, start=1):
            if place == 1:
                yield f'{place}. {member}: your turn'
            else:
                yield f'{place}. {member}: {place - 1} ahead of you'

    def __str__(self):
        return ' -> '.join(self) + ' -> None'


queue = BookingQueue('Python in Depth')
for member in ['Laura', 'Pablo', 'Marta', 'Iván']:
    queue.reserve(member)

print(queue.title, '|', len(queue), 'bookings')
print(queue)
for notice in queue.notices():
    print(' ', notice)
print('Is Marta in?', 'Marta' in queue, '| Is Sara in?', 'Sara' in queue)
print('Two walks:', list(queue), list(queue))
print('Book returned: it goes to', queue.serve())
print(queue)
print('Type of iter(queue):', type(iter(queue)).__name__)

The heart of the program is these two generators:

def __iter__(self):               # generator: walks the nodes without touching them
    node = self._first
    while node is not None:
        yield node.member
        node = node.next

def notices(self):
    for place, member in enumerate(self, start=1):
        if place == 1:
            yield f'{place}. {member}: your turn'
        else:
            yield f'{place}. {member}: {place - 1} ahead of you'

• node is a variable local to the generator, not an attribute of the queue. Each walk has its own, just like the separate cursors in the theory, and walking does not change the queue.
• notices() uses enumerate(self, start=1): enumerate asks __iter__ one at a time and puts the place in front. That is how Marta gets “2 ahead of you” without counting anything by hand.
• __str__ uses ' -> '.join(self): join accepts any iterable of strings, and the queue already is one.

In the drawing, the coral cursor node is on Pablo: it is the second next() of the walk. The queue has not noticed.

Output

Python in Depth | 4 bookings
Laura -> Pablo -> Marta -> Iván -> None
  1. Laura: your turn
  2. Pablo: 1 ahead of you
  3. Marta: 2 ahead of you
  4. Iván: 3 ahead of you
Is Marta in? True | Is Sara in? False
Two walks: ['Laura', 'Pablo', 'Marta', 'Iván'] ['Laura', 'Pablo', 'Marta', 'Iván']
Book returned: it goes to Laura
Pablo -> Marta -> Iván -> None
Type of iter(queue): generator

Let’s go through each line. Four bookings in order of arrival: Laura, Pablo, Marta, Iván. Notices are numbered from 1, and place p has p - 1 people ahead (Iván, 4 − 1 = 3). 'Marta' in queue gives True even though BookingQueue has no __contains__: if it is missing, Python walks the iterable with __iter__ and compares one by one. The two list(queue) calls in a row come out complete, because each one asked for a new generator. When the book comes back, serve() removes Laura and the queue starts at Pablo.

Generator step	node points to	Hands over	node moves to
1	Laura	'Laura'	Pablo
2	Pablo	'Pablo'	Marta
3	Marta	'Marta'	Iván
4	Iván	'Iván'	None
5	None	(leaves the while: StopIteration)	—

Questions you might ask

Why keep _last if _first already reaches everyone? So that reserve is instant. Without _last, each booking would have to walk the queue to the end (with 1000 members, 1000 jumps per booking).

Why does serve set _last = None when the queue becomes empty? Because otherwise _last would still point to the node we just removed, and the next booking would be linked to that “ghost” node instead of becoming the first one.

Why doesn’t __iter__ return self like the theory’s iterator? Because the queue is a container, not a walk. If it returned self with a cursor inside, print(queue) inside a loop over the queue would break the loop, just like the shelf’s 3 pairs.

Can I serve someone while walking the queue with a for? Better not. The generator follows the links at the moment it passes each node; if you change the structure halfway through, you may skip members or see someone who has already left. Walk first (for example, pending = list(queue)) and change afterwards.

Try it yourself

• Add cancel(member), which removes a member from any place (you will have to link the node in front to the one behind).
• Add backwards(), a generator that walks the queue from last to first. Hint: since the nodes only point forwards, walk once storing them in a list and then hand them over from the end to the beginning.

The pattern to take away

The three programs repeat three generator shapes you will recognise in any exercise:

Pattern	Skeleton	Where we used it
Infinite generator	while True: yield value; advance	fibonacci(), primes()
Generator that cuts	for x in other: if end: return; yield x	up_to(), primes_up_to(), n_primes()
__iter__ as a generator	node = first; while node: yield node.data; node = node.next	BookingQueue

Common mistakes in this practice

• Calling list() on an infinite generator: the program hangs with no error. Cut it earlier with up_to, islice or a break.
• Reusing the same generator for two questions: the second starts where the first ended. Create a new one for each walk.
• Writing the Fibonacci update in two lines: previous = current followed by current = previous + current uses the already-changed value. Use multiple assignment.
• Filtering an infinite generator with a comprehension: [p for p in primes() if p <= 150] never ends; the if does not stop the for.
• Storing the cursor as an attribute of a container: the second walk breaks the first. The cursor belongs in a local variable of the generator.
• Forgetting _last when the queue empties: the next booking gets lost, linked to a node that is no longer there.

Watch it step by step in Python Tutor

In Python Tutor, with Python 3, paste this reduced Fibonacci and step through it:

def fibonacci():
    previous, current = 0, 1
    while True:
        yield previous
        previous, current = current, previous + current


g = fibonacci()
a = next(g)
b = next(g)
c = next(g)
d = next(g)
e = next(g)
print(a, b, c, d, e)

Output:

0 1 1 2 3

What to watch:

1. After g = fibonacci() a generator object appears, but the fibonacci frame has not been created yet.
2. At a = next(g) the frame appears with previous = 0 and current = 1 and stops at the yield: a receives 0.
3. On each following next(), watch how previous and current change at the same time on the multiple-assignment line: (0, 1) → (1, 1) → (1, 2) → (2, 3).
4. Between one next() and the next, the generator frame stays there with its two variables: that is what “paused” means.

Summary and next step

• An infinite generator only keeps its state (two numbers in Fibonacci) and hands over values on demand.
• To cut it: return in another generator, islice, a counted next() or a for with break; never list() without a limit.
• Generators chain like a pipeline: each value travels the whole way before the next one is computed.
• return inside a generator ends it; the for walking it finishes with no error.
• An __iter__ written as a generator makes your own container work with for, list(), in, join and enumerate, and it can be walked as often as you like.
• A generator has no len(): count with sum(1 for _ in gen).

Checklist: have you practised it?

• Have you run the three programs and checked the outputs?
• Can you explain why previous, current = current, previous + current needs no temporary variable?
• Can you name three ways to cut an infinite generator?
• Do you understand why 'Marta' in queue works without __contains__?
• Have you tried the “Try it yourself” tasks without looking at the theory?

If you answered yes to almost everything, you are ready for the iterators and generators exercises, where I no longer give you the code. See you there!

Next step: exercisesIterators and generators in Python — exercises with multiples, a sequence that resets and perfect numbers (and a fine that came out letter by letter)

Did this article help you? Here is how you can support SergioLearns →