Python Iterators and Generators

An iterator gives you items one at a time, and a generator is the easiest way to make your own.

Iterables and Iterators

Anything you can loop over with for is an iterable: lists, tuples, strings, dictionaries, files and more.

An iterator is the object that actually walks through the items. It remembers where it is and gives the next item when asked.

Word	Meaning
Iterable	A collection you can loop over
Iterator	An object that returns one item at a time
`iter(x)`	Gets an iterator from an iterable
`next(it)`	Asks the iterator for the next item
`StopIteration`	The error raised when there are no items left

iter() and next()

Example

python

queue = ["Ali", "Sara", "Omar"]
it = iter(queue)

print(next(it))
print(next(it))
print(next(it))

try:
    next(it)
except StopIteration:
    print("No one left in the queue")

Output

Ali
Sara
Omar
No one left in the queue

This is exactly what a for loop does behind the scenes. It calls iter() once, then next() again and again, and stops quietly when it sees StopIteration.

Your Own Iterator Class

A class becomes an iterator when it has two methods: __iter__() and __next__().

Example

python

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

    def __iter__(self):
        return self

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

for n in Countdown(5):
    print(n, end=" ")
print("Go!")

Output

5 4 3 2 1 Go!

This works, but it needs a lot of code. Generators do the same job in a few lines.

Generators and yield

A generator is a function that uses yield instead of return. Each yield hands out one value and pauses the function. The next time you ask for a value, the function continues from where it stopped.

Example

python

def countdown(start):
    while start > 0:
        yield start
        start -= 1

for n in countdown(5):
    print(n, end=" ")
print("Go!")

gen = countdown(3)
print(next(gen))
print(next(gen))

Output

5 4 3 2 1 Go!
3
2

Same result as the class above, with a fraction of the code.

Why Generators Save Memory

A list keeps all its items in memory at once. A generator makes each item only when it is needed.

Example

python

import sys

big_list = [n * n for n in range(1_000_000)]
big_gen = (n * n for n in range(1_000_000))

print(sys.getsizeof(big_list) > 1_000_000)
print(sys.getsizeof(big_gen) < 500)
print(sum(big_gen))

Output

True
True
333332833333500000

The list takes megabytes. The generator stays tiny, yet sum() still gets every value.

A generator can be used only once. After it has given all its values, it is empty. If you need the values again, create a new generator or store them in a list.

A Practical Generator

Generators are great for reading big files or data line by line. Here is one that gives invoice numbers forever, and another that reads only the lines we care about.

Example

python

def invoice_numbers(prefix):
    n = 1
    while True:
        yield f"{prefix}-{n:04d}"
        n += 1

ids = invoice_numbers("KHI")
for _ in range(3):
    print(next(ids))

def paid_orders(lines):
    for line in lines:
        name, status = line.split(",")
        if status == "paid":
            yield name

orders = ["Ali,paid", "Sara,pending", "Omar,paid"]
print(list(paid_orders(orders)))

Output

KHI-0001
KHI-0002
KHI-0003
['Ali', 'Omar']

invoice_numbers() has an endless loop, but it never freezes the program, because it only runs when next() asks for a value.

Handy Tools in itertools

The itertools module has ready-made iterators:

Example

python

from itertools import count, islice, cycle

print(list(islice(count(100, 10), 4)))

turns = cycle(["Ali", "Sara"])
print([next(turns) for _ in range(5)])

Output

[100, 110, 120, 130]
['Ali', 'Sara', 'Ali', 'Sara', 'Ali']

islice() takes only the first few items from an endless iterator.

Exercise

What keyword turns a normal function into a generator?

returngenerateyieldnext