Python Generator Cheat Sheet

Reference for Python generator syntax. For a detailed introduction to generators, see blog posts on iterators and generators.

Source

Find source code on GitHub.

Background

Generators can be used to make a class or a function iterable. Python will create a generator when it encounters the yield keyword in a function.

To make a class iterable, implement an __iter__() method with yield. An instance of the class can then be passed directly into a for loop.

class Iterable:
    def __iter__(self):
        # loop over collection
        yield   # each item in collection

for item in Iterable():
    # use item

__iter__() will return a generator object that is only evaluated when necessary.

To make a function iterable, use yield:

def function():
    # loop over collection
    yield   # each item in collection

for item in function():
    # use item

Examples

Iterable class

Implement __iter__() with yield:

class Iterable:
    def __init__(self, collection):
        self.collection = collection

    def __iter__(self):
        for item in self.collection:
            yield item

Then pass an instance of the class to for:

>>> collection = [7, 8, 9]
>>> iterable = Iterable(collection)
>>> for item in iterable:
...     print(item)
...
7
8
9

Iterable function

Implement a function with yield:

def iterate_over(collection):
    for item in collection:
        yield item

Then call the function in for:

>>> collection = [7, 8, 9]
>>> for item in iterate_over(collection):
...     print(item)
...
7
8
9

Return

A generator iterator raises the StopIteration exception when it returns, either at the end of a function or from the return keyword.

def iterable_function():
    for i in range(1, 10):
        if i > 3:
            return
        yield i
>>> for item in iterable_function():
...     print(item)
...
1
2
3

Further Reading

Links to Python documentation in no particular order.