The Invent with Python Blog

Fri 24 September 2021

What is a Python Generator? (Implementing Your Own range() Function)

Posted by Al Sweigart in misc   

Generators in Python (also called generator functions) are used to create a series of values one at a time. This can keep your program from requiring the large amounts of memory needed if you generated all the values in the series at once. For example, in Python 2 running for i in range(1000000): the range() function would create and return a list with one million integers in it. This takes up a large amount of memory even though the for loop only needs one integer at a time. This was fixed in Python 3, where range() now produces one integer at a time. Generator functions let you create one value at a time with any sort of data, not just ranges of integers.

If you find yourself creating a large list of values with, say, a list comprehension or for loop, but you only need one value at a time, you should consider using a generator instead. This is especially true if you are creating large lists with hundreds of thousands of items or more. Using a generator is like using the readline() method to read a text file one line at a time, instead of using the read() method to read in the entire file all at once.

As a side note, generator functions return generator objects and generator objects are iterables. (Iterables are beyond the scope of this blog post, but see The Iterator Protocol: How "For Loops" Work in Python for more details.)

You can download a single .py file of all the examples from whatIsAGenerator.py. This blog post assumes you have a basic beginner's level of understanding of Python.

First, let's take a look at a regular function that you're already familiar with:

### A regular function:
def aRegularFunction(param):
    print('Hello', param)
    return 42

print("Calling aRegularFunction('Alice'):")
returnedValue = aRegularFunction('Alice')
print('Returned value:', returnedValue)

When you run this code, the output looks like this:

Calling aRegularFunction('Alice'):
Hello Alice
Returned value: 42

This is familiar and unsurprising. A function takes arguments for its parameters, runs some code, and then returns some return value.

Now let's take a look at a generator function. You can tell a function is a generator function because it'll have the yield keyword somewhere in the function body:

### A generator (aka generator function):
def aGeneratorFunction(param):
    print('Hello', param)
    yield 42
    print('How are you,', param)
    yield 99
    print('Goodbye', param)
    return 86  # Raises StopIteration with value 86.

print("Calling aGeneratorFunction('Bob'):")
generatorObj = aGeneratorFunction('Bob')  # Does not run the code, but returns a generator object.
print('Calling next():')
yieldedValue = next(generatorObj)
print('Yielded value:', yieldedValue)
print('Calling next():')
yieldedValue = next(generatorObj)
print('Yielded value:', yieldedValue)
print('Calling next():')
try:
    next(generatorObj)
except StopIteration as excObj:
    print('type(excObj) is', type(excObj))
    print('excObj.value is', excObj.value)
    print('type(excObj.value) is', type(excObj.value))

When you run this code, the output looks like this:

Calling aGeneratorFunction('Bob'):
Calling next():
Hello Bob
Yielded value: 42
Calling next():
How are you, Bob
Yielded value: 99
Calling next():
Goodbye Bob
type(excObj) is <class 'StopIteration'>
excObj.value is 86
type(excObj.value) is <class 'int'>

What you need to know is that one does not simply call a generator function. Calling a generator function does not run it's code. Instead, calling a generator function returns a generator object. You can see this in the interactive shell:

>>> aGeneratorFunction('Eve')
<generator object aGeneratorFunction at 0x0000013470179580>

To run the code inside a generator function, you must call Python's built-in next() function and pass it the generator object. This will run the code up until a yield statement, which acts like a return statement and makes the next() call return a value. The next time next() is called, the execution resumes from the yield statement with all the same values for the local variables.

When a return statement is reached, the generator function raises a StopIteration exception. Usually there's no value returned, but if so it is set to the exception object's value attribute.

The next() function usually isn't used with generators, but rather the generator function is used in a for loop. On each iteration, the for loop sets the loop variable to the yielded value. The StopIteration exception tells the for loop that the generator object has exhausted its values. For example:

### Using a generator function with a for loop or list() call:
for yieldedValue in aGeneratorFunction('Carol'):
    print('Yielded value:', yieldedValue)

listOfGeneratedValues = list(aGeneratorFunction('David'))
print('List of generated values:', listOfGeneratedValues)

When you run this code, the output looks like this:

Hello Carol
Yielded value: 42
How are you, Carol
Yielded value: 99
Goodbye Carol
Hello David
How are you, David
Goodbye David
List of generated values: [42, 99]

Rememebr that generators are mainly used to generate a series of values, but one at a time. Let's use an example of a Fibonacci sequence number generator. The Fibonacci sequence begins with 0 and 1, and then the next number in the sequence is the sum of the previous two numbers: 0, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55, and so on.

A generator can produce these numbers (up to a given maximum) one at a time. This generator can be used with a for loop or a list() call:

### A practical example of a generator using the Fibonacci sequence:
def fibonacciSequence(maxNum=50):
    a = 0
    b = 1
    yield a
    yield b
    while True:
        a, b = b, b + a
        if b >= maxNum:
            return b  # Raise StopIteration.
        else:
            yield b


for fibNum in fibonacciSequence():
    print('Next number in the Fibonacci sequence:', fibNum)


fibNumsUpTo500 = list(fibonacciSequence(500))
print('List of fib numbers up to 500:', fibNumsUpTo500)

When you run this code, the output looks like this:

Next number in the Fibonacci sequence: 0
Next number in the Fibonacci sequence: 1
Next number in the Fibonacci sequence: 1
Next number in the Fibonacci sequence: 2
Next number in the Fibonacci sequence: 3
Next number in the Fibonacci sequence: 5
Next number in the Fibonacci sequence: 8
Next number in the Fibonacci sequence: 13
Next number in the Fibonacci sequence: 21
Next number in the Fibonacci sequence: 34
List of fib numbers up to 500: [0, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 377]

To summarize, Python generators look like functions that contain the yield keyword. They're used for producing a series of values one at a time so that your program doesn't consume a lot of memory by producing all the values at once. When a generator function is called, it doesn't run the code inside it but rather returns a generator object. This generator object can be passed to the built-in next() function. This runs the code in the generator function up until a yield statement (which gives a value for next() to return) or a return statement (which raises the StopIteration to signal that there are no more values to generate). However, generators are usually used in for loops, which automatically handle them.

That's it! Generators aren't all that complicated. I recommend you learn about iterables and iterators in general by checking out Trey Hunner's blog post, The Iterator Protocol: How "For Loops" Work in Python.

Implementing Your Own range() Function

Let's put our knowledge to the test by creating our own range() function. Instead of for i in range(10):, we'll be able to use for i in mySimpleImplementationOfRange(10):. We'll use a generator function that takes a single integer argument, and yields values starting from 0 up to but not including the argument:

### Using generator functions to re-implement range():
def mySimpleImplementationOfRange(stop):
    i = 0
    while i < stop:
        yield i
        i += 1
    return  # Raise StopIteration.

for i in mySimpleImplementationOfRange(5):
    print(i)

When you run this code, the output looks like this:

0
1
2
3
4

Of course, the range() function is a bit more complicated than that. It can take up to three arguments to specify a starting integer, stopping integer, and a "step" integer to indicate how much the yielded integers should change by. Also, range() only accepts integer arguments and the step argument cannot be 0. A fuller implementation is below:

### Using generator functions for a more complete range() re-implementation:
def myImplementationOfRange(firstParam, secondParam=None, thirdParam=None):
    if secondParam is None:
        if not isinstance(firstParam, int):
            raise TypeError("'" + firstParam.__class__.__name__ + "' object cannot be interpreted as an integer")
        start = 0
        stop = firstParam
        step = 1
    elif thirdParam is None:
        if not isinstance(secondParam, int):
            raise TypeError("'" + secondParam.__class__.__name__ + "' object cannot be interpreted as an integer")
        start = firstParam
        stop = secondParam
        step = 1
    else:
        if not isinstance(thirdParam, int):
            raise TypeError("'" + thirdParam.__class__.__name__ + "' object cannot be interpreted as an integer")
        if thirdParam == 0:
            raise ValueError('myImplementationOfRange() arg 3 must not be zero')
        start = firstParam
        stop = secondParam
        step = thirdParam

    i = start
    if step > 0:  # step arg is increasing
        while i < stop:
            yield i
            i += step
        return i
    elif step < 0:  # step arg is decreasing
        while i > stop:
            yield i
            i += step  # Adding a negative to decrease
        return i

for i in myImplementationOfRange(4, 12, 2):
    print(i)

When you run this code, the output looks like this:

4
6
8
10

There's a bit more to it than that, as range is actually a class which implements several dunder methods (which you can learn about from Nina Zakharenko's PyCon 2018 talk, Elegant Solutions for Everyday Python Problems. But now you can see how something like Python's built-in range() is implemented.

 

Learn to program with my books for beginners, free under a Creative Commons license:

Take my Automate the Boring Stuff with Python online Udemy course. Use this link to apply a 60% discount.