Python - Generators

The task of implementing iterators can be simplified by using generators. We have seen how to create custom iterators using the object-oriented way, i.e., by defining a class that has __init__, __next__, and __iter__ methods. For example, we saw the Cubes class which when instantiated created an iterator object that gave out cubes of numbers. These types of simple iterators can be implemented in a much easier way by writing generators.

There are two types of generators: generator functions and generator expressions. Both of them are used to create generator objects which are actually iterators. A generator object is a kind of iterator, and we get a generator object by writing a generator function or a generator expression.

When you write a generator, you do not need to worry about writing the __iter__ and __next__ methods. You get the iterator interface automatically. So, when you want to get an iterator without writing a class, you can write generators. In fact, writing a class to define your own iterator is very rare. Generally, the automated syntax of generators is preferred to get your own iterators. However, if you need to create complex iterators or need to give access to some extra attributes and methods, then you will have to write class-based iterators.

>>> def cubes(start, stop):

...         for n in range(start, stop+1):

...             yield n * n * n

>>> y = cubes(2, 5)

Right now, do not worry about what is written inside the function. What you need to understand is that when we call this function, it returns a generator object, which is actually an iterator.

>>> y

<generator object cubes at 0x000001E3263A75B0>

>>> y is iter(y)

True

By writing this generator function, we were able to get an iterator without worrying about any of the methods needed to satisfy the iterator protocol. They are automatically implemented for us.

We can use the next function to get values from this generator object.

>>> next(y)

8

>>> next(y)

27

>>> next(y)

64

>>> next(y)

125

>>> next(y)

StopIteration

The calls to next function give us the cubes, and the StopIteration error was raised to signify the end of data. This generator object y behaves just like the iterator object x (instance of Cubes) would have behaved.

Generator expression is an expression that returns an iterator also known as generator object. This generator object returns values one by one when used in an iteration context such as for loop.

Generator expressions are syntactically almost similar to list comprehensions, but the difference is that generator expressions return a generator object instead of a list. Generator objects generate one value at a time, while the comprehensions save all the values in memory. This is why generator expressions consume less memory as compared to comprehensions, and there is no waiting time as all the values are not computed at once. The saving in memory and time is crucial when the number of data values is very large. Let us see an example:

>>> (n * n * n for n in range(2, 6))

<generator object <genexpr> at 0x000001F5FD451E70>

This is an example of a generator expression. We know that list comprehensions are enclosed in square brackets, while set and dictionary comprehensions are enclosed in curly braces and we do not have anything like tuple comprehensions, so when we have a comprehension like expression enclosed in parentheses, it is a generator expression. This generator expression gives us a generator object which is lazily evaluated, and it can be iterated over.

>>> g = (n * n * n for n in range(2, 6))

>>> g

<generator object <genexpr> at 0x000001F5FF58B1B0>

g is a generator object, let us call next function for it.

>>> next(g)

8

>>> next(g)

27

>>> next(g)

64

>>> next(g)

125

>>> next(g)

StopIteration

Now it is exhausted, so we cannot use it again.

>>> for i in (n * n * n for n in range(2, 6) if n % 2 == 0):

...  print(i)

Generator expressions are written inside parentheses. If you are writing the generator expression as a single argument to a function call, then the parentheses of the function call are sufficient, there is no need to write 2 pairs of parentheses, one for the call and other for the generator expression. But if there are more than one argument, then you need to enclose the generator expression in parentheses otherwise you will get a syntax error.

For example, in the following function call, we have sent a generator expression as argument.

>>> func(n * n for n in range(2,4))

Since this is the only argument, the parentheses are not required. If we have more arguments, then we will have to put the parentheses.

func((n*n for n in range(2,4)), 'x')