Python Notes and Examples

Iterators and Generators

Iterators

The for loop iterates over an iterable (such as a list) by implicitly making use of the iterable’s implementing the iterator protocol. Here’s how you can manually do it:

x = range(4)
type(x) #=> <class 'range'>

y = iter(x)
type(y) # an iterator object.

next(y) #=> 0
next(y) #=> 1
next(y) #=> 2
next(y) #=> 3
next(y) # Throws a StopIteration exception.

If you want to be able to loop over instances of one of your own classes (the same way you can loop over lists, strings, dicts, and sets), then you’ll need to have your class implement the iterator protocol. To do this, your class must implement the __iter__() method. The __iter__() method should return an iterator — an object with a __next__() method — such that next(the_iterator) works on it.

Note: strings, lists, dicts, sets, and files all implement the iterator protocol. They are not iterators themselves:

>>> z = [11, 12, 13]
>>> next(z)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: 'list' object is not an iterator

When you try to loop over anything, Python implicitly calls iter(the_object) behind the scenes for you.

Generators

Functions that contain a yield expression are so-called “generator functions”. When you call them they return a generator object. A generator object implements the iterator protocol:

>>> def f():
...  yield 10
...  for i in range(3):
...   yield i
...  yield 100
...
>>> x = f()  # `x` is a generator object.
>>> next(x)
10
>>> next(x)
0
>>> next(x)
1
>>> next(x)
2
>>> next(x)
100
>>> next(x)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

>>> y = f()
>>> for i in y:
...  print(i)
...
10
0
1
2
100

# Eagerly get all values from the iterator:
z = list(f())

You might use a generator in situations where you only need elements as-needed in a list that’s time-consuming to generate:

def get_computation_intensive_results(coll):
    for x in coll:
        yield time_consuming_computation(x)

Finally, note that there’s a shorthand for creating generators: it’s like a list comprehension but with parens instead of brackets.

foo = (x**2 for x in range(5))
next(foo) #=> 0
next(foo) #=> 1
next(foo) #=> 2
# etc.

some_func((x**2 for x in range(5)))
some_func( x**2 for x in range(5) ) # Can omit the extra parens.

If, in your generator function, you’re yielding values from iterables, you can can use the “yield from” syntax to save a few lines of code:

def gen1(n):
    for i in get_some():
        yield i
    for i in get_more():
        yield i

def gen2(n):
    yield from get_some()
    yield from get_more()

Note:

After the generator object yields a value via its generator function’s yield expression, the code calling the generator can send a value back into the generator object — becoming the value of the yield expression for possible use in the next value generated by the generator. This is how Python implements coroutines; see PEP 342.