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

While preparing the library fines for a class assignment I wrote a function that, for each overdue loan, worked out the fine in cents. I ran list(fines(loans)) expecting two fines and got this: ['P', 'a', 'b', 'l', 'o', ':', ' ', '8', '0', ' ', 'c', 't', 's']. One fine, split into letters. I had written return where I wanted yield. That mistake, and the other two that keep coming up with generators, are in the “What’s wrong here?” section of this article.

You have already seen the theory and the practice. Now it is your turn, with three exercises on three levels, three tricky snippets, an exam-style quiz and the cheat sheet for the topic.

How to use this article:

1. Read the statement and the expected output of each exercise, and write your solution before looking at anything else.
2. If you have been stuck for more than 10 minutes, open the hint.
3. Before the solutions, try to predict what the “What’s wrong here?” snippets do.
4. Compare your code with the commented solution: what matters is understanding why it is done that way.

Python iterators and generators exercises — Basic Level

Exercise 1 — Generator of multiples of a number

Write three generators of multiples:

• infinite_multiples(base): base, 2 × base, 3 × base… without stopping.
• multiples_up_to(base, limit): the multiples that do not exceed the limit (the limit included).
• first_n_multiples(base, n): the first n, using infinite_multiples inside.

With them (and a generator expression with if or two), produce these four lists. The output should be:

Multiples of 3 up to 30:
3 6 9 12 15 18 21 24 27 30

First 8 multiples of 7:
7 14 21 28 35 42 49 56

Multiples of 5 between 20 and 50:
20 25 30 35 40 45 50

Common multiples of 3 and 5 up to 100:
15 30 45 60 75 90

Python iterators and generators exercises — Intermediate Level

Exercise 2 — Arithmetic sequence iterator

Implement the class ArithmeticSequence(begin, step, terms) as an iterator written by hand (with __iter__ and __next__, no yield). For example, ArithmeticSequence(2, 3, 5) gives 2, 5, 8, 11, 14: it starts at 2, adds 3 and gives 5 terms.

Requirements:

• __init__ raises ValueError('The step cannot be 0') if the step is 0, and another ValueError if the terms are not a positive integer.
• It must be walkable several times: with for, with list(), with sum()… one after another, all of them complete.
• __len__ returns the terms still left in the current walk (after two next() calls on a 5-term sequence, 3 are left).
• __str__ shows the whole sequence with arrows: 2 → 5 → 8 → 11 → 14.

The output should be:

Sequence(2, step=3, terms=5):
  for:   2 5 8 11 14
  list:  [2, 5, 8, 11, 14]
  sum:   40
  str:   2 → 5 → 8 → 11 → 14
  len after two next(): 3
Sequence(10, step=-2, terms=6):
  for:   10 8 6 4 2 0
Sequence(0, step=0.5, terms=8):
  for:   0 0.5 1.0 1.5 2.0 2.5 3.0 3.5
Expected error: The step cannot be 0

Python iterators and generators exercises — Final Challenge

Exercise 3 — Perfect number generator

A number is perfect if it equals the sum of its proper divisors (all of them except itself). 6 is: 1 + 2 + 3 = 6. So is 28: 1 + 2 + 4 + 7 + 14 = 28. Write:

• proper_divisors(n): the sorted list of proper divisors. Make it fast: there is no need to try every number up to n.
• is_perfect(n).
• perfect_numbers(): an infinite generator of perfect numbers.
• perfect_up_to(limit): a generator with the perfect numbers that do not exceed the limit, which ends even if there are no more perfect numbers nearby.

The output should be:

Is 6 perfect? True
Is 12 perfect? False
Is 28 perfect? True

Perfect numbers up to 10000:
6 28 496 8128

Divisors of 6: [1, 2, 3] → sum: 6
Divisors of 28: [1, 2, 4, 7, 14] → sum: 28
Divisors of 496: [1, 2, 4, 8, 16, 31, 62, 124, 248] → sum: 496

What’s wrong here?

Three snippets with the same overdue-loans generator (15-day loan period). Before reading the explanation, predict what each one prints.

Snippet 1: the report that stops halfway

def overdue(loans, max_days=15):
    for member, days in loans:
        if days > max_days:
            yield member


loans = [('Pablo', 19), ('Laura', 3), ('Marta', 22)]
notices = overdue(loans)

print('There are', len(list(notices)), 'notices')
for member in notices:
    print('Notice for', member)
print('End of report')

What does it print? Think before reading on.

There are 2 notices
End of report

It says there are 2 notices and prints none. len(list(notices)) walked the whole generator to count it; when the for arrives, the generator is already at the end and runs zero laps. There is no error because, for Python, walking an exhausted generator is perfectly valid: there is simply nothing there.

The fix: if you will use it more than once, store it in a list from the start.

def overdue(loans, max_days=15):
    for member, days in loans:
        if days > max_days:
            yield member


loans = [('Pablo', 19), ('Laura', 3), ('Marta', 22)]
notices = list(overdue(loans))   # a list CAN be walked several times

print('There are', len(notices), 'notices')
for member in notices:
    print('Notice for', member)
print('End of report')

Output:

There are 2 notices
Notice for Pablo
Notice for Marta
End of report

Snippet 2: the generator without brackets

def overdue(loans, max_days=15):
    for member, days in loans:
        if days > max_days:
            yield member


loans = [('Pablo', 19), ('Laura', 3), ('Marta', 22)]
for member in overdue:
    print('Notice for', member)

What does it print?

Traceback (most recent call last):
  File "not_called.py", line 8, in <module>
    for member in overdue:
                  ^^^^^^^
TypeError: 'function' object is not iterable

overdue is the function; what can be walked is what it returns when called, overdue(loans). Without the brackets we are asking the for to walk a function, and a function has no __iter__. The message says it exactly: 'function' object is not iterable.

Fix: for member in overdue(loans):.

Notice for Pablo
Notice for Marta

Snippet 3: the fine letter by letter

def fines(loans, max_days=15):
    for member, days in loans:
        if days > max_days:
            return f'{member}: {(days - max_days) * 20} cts'


loans = [('Pablo', 19), ('Laura', 3), ('Marta', 22)]
print(list(fines(loans)))

What does it print?

['P', 'a', 'b', 'l', 'o', ':', ' ', '8', '0', ' ', 'c', 't', 's']

It is the story from the beginning. Since fines has return and no yield, it is not a generator: it is a normal function that, at the first overdue loan (Pablo), returns the string 'Pablo: 80 cts' and ends. Marta is never looked at. And list() on a string splits it into characters, because a string is iterable too. The numbers add up: (19 − 15) × 20 = 80 cents.

With yield instead of return, the function becomes a generator and hands over one fine per overdue loan: Pablo, (19 − 15) × 20 = 80, and Marta, (22 − 15) × 20 = 140.

['Pablo: 80 cts', 'Marta: 140 cts']

Commented solutions

Solution Exercise 1

def infinite_multiples(base):
    """Infinite generator: base, 2*base, 3*base..."""
    multiple = base
    while True:
        yield multiple
        multiple += base         # adding is cheaper than multiplying every time


def multiples_up_to(base, limit):
    """Finite generator: stops by itself once past the limit."""
    multiple = base
    while multiple <= limit:     # <= because the limit counts too (30 is in)
        yield multiple
        multiple += base


def first_n_multiples(base, n):
    counter = 0
    for m in infinite_multiples(base):
        if counter == n:
            return               # we cut the infinite one with return
        yield m
        counter += 1


def show(title, gen):
    print(title)
    print(' '.join(str(m) for m in gen))   # consumes the whole generator once


show('Multiples of 3 up to 30:', multiples_up_to(3, 30))
show('\nFirst 8 multiples of 7:', first_n_multiples(7, 8))
show('\nMultiples of 5 between 20 and 50:', (m for m in multiples_up_to(5, 50) if m >= 20))
show('\nCommon multiples of 3 and 5 up to 100:', (m for m in multiples_up_to(3, 100) if m % 5 == 0))

Why it is done this way:

• multiples_up_to uses while multiple <= limit (not <) because the statement includes the limit: 30 has to come out in the multiples of 3 up to 30.
• Adding base on each lap is cheaper and clearer than multiplying base * k with a separate counter.
• first_n_multiples cuts an infinite generator with a counter and return: it is the “generator that cuts” pattern from the practice. With counter == n checked before the yield, not a single extra multiple is computed.
• The last two cases need no new functions: a generator expression with if filters what comes out of multiples_up_to.
• show receives any generator and consumes it once with ' '.join(...). That is why each call gets a new one.

Solution Exercise 2

class ArithmeticSequence:
    def __init__(self, begin, step, terms):
        if step == 0:
            raise ValueError('The step cannot be 0')
        if not isinstance(terms, int) or terms <= 0:
            raise ValueError('The terms must be a positive integer')
        self.__begin = begin          # the original data never changes...
        self.__step = step
        self.__terms = terms
        self.__current = begin        # ...and the walk state is reset in __iter__
        self.__remaining = terms

    def __iter__(self):
        self.__current = self.__begin
        self.__remaining = self.__terms
        return self

    def __next__(self):
        if self.__remaining == 0:
            raise StopIteration
        value = self.__current
        self.__current += self.__step
        self.__remaining -= 1
        return value

    def __len__(self):
        return self.__remaining       # the ones STILL to come in the current walk

    def __str__(self):
        return ' → '.join(str(self.__begin + k * self.__step) for k in range(self.__terms))


def try_seq(begin, step, terms, full=False):
    seq = ArithmeticSequence(begin, step, terms)
    print(f'Sequence({begin}, step={step}, terms={terms}):')
    print('  for:  ', ' '.join(str(v) for v in seq))
    if full:
        print('  list: ', list(seq))  # works again because __iter__ resets
        print('  sum:  ', sum(seq))
        print('  str:  ', seq)
        it = iter(seq)
        next(it)
        next(it)
        print('  len after two next():', len(seq))


try_seq(2, 3, 5, full=True)
try_seq(10, -2, 6)
try_seq(0, 0.5, 8)
try:
    ArithmeticSequence(1, 0, 5)
except ValueError as err:
    print('Expected error:', err)

What makes it walkable several times is in these lines:

def __iter__(self):
    self.__current = self.__begin
    self.__remaining = self.__terms
    return self

def __next__(self):
    if self.__remaining == 0:
        raise StopIteration
    value = self.__current
    self.__current += self.__step
    self.__remaining -= 1
    return value

def __len__(self):
    return self.__remaining       # the ones STILL to come in the current walk

• __init__ stores the original data (__begin, __step, __terms) and also prepares the state, in case someone calls next() without going through iter().
• __iter__ puts __current and __remaining back at the start and returns self. Since for, list() and sum() call iter() before starting, each walk begins at the first term.
• __next__ hands over the current value and then moves on; when __remaining reaches 0, it raises StopIteration.
• __len__ returns __remaining: after iter() and two next() calls on 5 terms, 5 − 2 = 3 are left.

Trace of ArithmeticSequence(2, 3, 5) in one walk:

The sum: 2 + 5 + 8 + 11 + 14 = 40. And __str__ does not use the iterator: it computes each term with begin + k * step, so printing the sequence does not spoil a walk that is halfway through.

Solution Exercise 3

def proper_divisors(n):
    """Divisors of n except n itself, in order. Only goes up to the square root."""
    if n < 2:
        return []
    small, big = [1], []
    i = 2
    while i * i <= n:
        if n % i == 0:
            small.append(i)    # i divides n...
            if i != n // i:    # ...and so does its partner n // i (root not repeated)
                big.append(n // i)
        i += 1
    return small + big[::-1]   # partners come out largest first: reverse them


def is_perfect(n):
    return n >= 2 and sum(proper_divisors(n)) == n


def perfect_numbers():
    """Infinite generator of perfect numbers."""
    candidate = 2
    while True:
        if is_perfect(candidate):
            yield candidate
        candidate += 1


def perfect_up_to(limit):
    """Checks the CANDIDATES up to the limit, not the perfect ones: so it ends even if there are no more."""
    for candidate in range(2, limit + 1):
        if is_perfect(candidate):
            yield candidate


for n in (6, 12, 28):
    print(f'Is {n} perfect? {is_perfect(n)}')

print('\nPerfect numbers up to 10000:')
print(' '.join(str(p) for p in perfect_up_to(10000)))

print()
for p in perfect_up_to(500):
    divs = proper_divisors(p)
    print(f'Divisors of {p}: {divs} → sum: {sum(divs)}')

proper_divisors looks for divisors in pairs:

def proper_divisors(n):
    """Divisors of n except n itself, in order. Only goes up to the square root."""
    if n < 2:
        return []
    small, big = [1], []
    i = 2
    while i * i <= n:
        if n % i == 0:
            small.append(i)    # i divides n...
            if i != n // i:    # ...and so does its partner n // i (root not repeated)
                big.append(n // i)
        i += 1
    return small + big[::-1]   # partners come out largest first: reverse them

For 28 it only tries i = 2, 3, 4 and 5 (6 × 6 = 36 is already over): 2 brings its partner 14 and 4 brings 7. The condition i != n // i avoids adding the root twice when the number is a square (in 16, the 4 would be repeated). The big ones come out reversed (14, 7), which is why they are flipped with [::-1] before joining them with the small ones.

And here is the important decision in the exercise:

def perfect_up_to(limit):
    """Checks the CANDIDATES up to the limit, not the perfect ones: so it ends even if there are no more."""
    for candidate in range(2, limit + 1):
        if is_perfect(candidate):
            yield candidate

perfect_up_to walks the candidates from 2 to the limit and stops when it runs out of them, whatever it has found. The version that looks natural, for p in perfect_numbers(): if p > limit: return, only checks the limit when a new perfect number appears, and the one after 8128 is 33,550,336: it would have to try thirty-three million candidates before it could stop. The old version of this article had exactly that bug.

Exam-style quiz

Now test yourself as in the lab. FP2 quizzes are like this: about 10 questions in 10-15 minutes, single answer, multiple answer (wrong answers subtract), true or false, fill in the gaps, matching and “what does this code print?”. Here are 13 original questions on iterators and generators. Answer without running anything, press Check and read the explanation of each one. These questions are also the bank for the FP2 theory mock exam, which will mix them with those of the other topics.

def evens(n):
    for i in range(n):
        if i % 2 == 0:
            yield i

g = evens(7)
next(g)
print(list(g))

squares = (x * x for x in range(4))
print(sum(squares), sum(squares))

def gen():
    print("A", end="")
    yield 1
    print("B", end="")
    yield 2
    print("C", end="")

g = gen()
print("0", end="")
for x in g:
    print(x, end="")
print()

class Countdown:
    def __init__(self, start):
        self.n = start

    def __iter__(self):
        return self

    def __next__(self):
        if self.n <= 0:
            raise StopIteration
        self.n -= 1
        return self.n + 1

c = Countdown(3)
print(list(c), list(c))

def everything(*lists):
    for l in lists:
        yield from l

print(list(everything([1, 2], [], [3])))

it = iter(["a", "b"])
next(it)
next(it)
print(next(it, "end"))

Watch it step by step in Python Tutor

Let’s watch snippet 1’s bug live. In Python Tutor, choose Python 3, paste this and step through it with Next:

def overdue(loans):
    for member, days in loans:
        if days > 15:
            yield member


loans = [('Pablo', 19), ('Laura', 3), ('Marta', 22)]
notices = overdue(loans)
first_run = list(notices)
second_run = list(notices)
print(first_run, second_run)

Output:

['Pablo', 'Marta'] []

What to watch:

1. notices is a generator object: the function has not started yet.
2. During first_run = list(notices) the overdue frame appears and disappears: it hands over Pablo, skips Laura, hands over Marta and ends.
3. At second_run = list(notices) the arrow does not even enter the function: the generator has already finished and list() gets StopIteration straight away.
4. At the end, first_run is ['Pablo', 'Marta'] and second_run is []. Same object, two results.

Cheat sheet: iterators and generators in Python

The essentials in 30 seconds: an iterable can be walked (it has __iter__); an iterator walks (it has __next__ and returns itself in __iter__); the for does iter() + next() until StopIteration. A function with yield returns a generator: lazy, pausable and single-use.

Templates from the exercises:

Before handing in, ask yourself:

1. Will this generator be walked more than once? If so, do I create it again or keep a list?
2. Does every infinite generator have something to cut it (return, break, islice)?
3. Does my __iter__ reset the state or, better, return a new generator?
4. Have I written yield where I want one value per lap and return only where I want to finish?

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