How to Write a Python Generator

Iterators play an important role in Python. They allow you to create an object that can be iterated upon. It’s a good way to store data that can be accessed through a for in loop.

The trouble with iterators is that it takes a lot of work to build one. They are useful, but you have to write a lot of code to make one work. Generators help solve this problem.

In this guide, we’re going to talk about what Python generators are and why you should use them. We’ll also implement a generator in Python to help you understand how they work.

Without further ado, let’s get started!

What is a Generator?

A generator is a simple way of creating an iterator in Python. It is a function that returns an object over which you can iterate.

Generators are often called syntactic sugar. This is because they do not necessarily add new functionality into Python. They help make your code more efficient.

Like iterators, you can only iterate over a generator once. This is because generators keep track of how many times they have been iterated upon, and cannot be reset.

Let’s take a look at a Python iterator:

class Cookies:
	def __init__(self, value):
		self.value = value

	def __iter__(self):
		return self

	def __next__(self):
		return self.value

This code creates an iterator called Cookies. When you pass data through this iterator, it will be turned into an iterable object. This means that you can loop through the data using a for in loop. If you’re looking at this code and think it is really long, that’s because it is.

Generators can help us shorten this code. Let’s write a generator for our above iterator:

def cookies(value):
	while True:
		yield value

Our code has been drastically reduced in size. Let’s discuss how this works.

A generator is a function that has a yield keyword instead of a return statement. Yield statements return a value from a function.

The difference between a yield statement and a return statement is that a return statement stops a function from running, whereas a yield statement pauses the function and continues on iterating.

Let’s try to iterate over our simple generator function:

for cookie in cookies("Raspberry"):
	print(cookie)

Our code returns:

Raspberry

Raspberry

…

This code keeps repeating until we stop our program. This is because we have used a while loop which executes forever. This isn’t very useful in most cases. Let’s write a generator that stops when an action has been performed.

How to Write a Python Generator

We have a list of cookies that we want to print to the console. This list looks like this:

[“Raspberry”, “Choc-Chip”, “Cinnamon”, “Oat”]


To print these out to the console, we could create a simple generator. Open up a new Python file and paste in the following code:

def print_cookies(cookies):
	length = len(cookies)
	for cookie in range(0, length):
yield cookies[cookie]

Our generator goes through every cookie in the list we specify and returns each cookie individually. This generator does not work just yet. We have to use a for loop to iterate over it:

cookie_list = ["Raspberry", "Choc-Chip", "Cinnamon", "Oat"]

for c in print_cookies(cookie_list):
	print(c)

We’ve defined an array called cookie_list which stores a list of four cookies. We have then set up a for loop which uses our generator to iterate through the values in cookie_list.

For each iteration in our for loop, the generated object is printed to the console:

Raspberry

Choc-Chip

"Career Karma entered my life when I needed it most and quickly helped me match with a bootcamp. Two months after graduating, I found my dream job that aligned with my values and goals in life!"

Venus, Software Engineer at Rockbot

Find Your Bootcamp Match

Cinnamon

Oat

We did it! This generator returns values until each value in the variable cookie_list has been iterated over.

Generators have a range of use cases. Let’s say that we have a list of order prices that we want to total up. We could do this using the following code:

def total_orders(orders):
	length = len(orders)
	total = 0
	for o in range(0, length):
		total += orders[o]
		yield total

This generator will keep a running total of the value of all orders in a list. After each item has been iterated over, the generator will yield the current total value of the orders. Let’s write a for loop which uses our generator:

order_list = [2.30, 2.50, 1.95, 6.00, 7.50, 2.15]

for order in total_orders(order_list):
	print(order)

Our code returns:

2.3

4.8

6.75

12.75

20.25

22.4

Our generator has calculated the cumulative value of all the orders in the order_list variable. Every iteration over the list returns the new cumulative value.

How to Write a Generator Expression

Generators are already easier to write than an iterator. Our quest for writing cleaner code does not need to stop, thanks to generator expressions.

Generator expressions are similar to list comprehensions. Generator expressions produce one item at a time, like a generator. This is different from a list comprehension which produces an entire list, all at one time.

Let’s write a generator for our cookies example:

cookie_list = ["Raspberry", "Choc-Chip", "Cinnamon", "Oat"]

cookie_generator = (cookie for cookie in cookie_list)

We’ve defined a list of cookies in the variable cookie_list. We then create a generator expression. This expression uses the syntax for list comprehensions but with one big difference: list comprehensions are defined within square brackets, whereas generators are defined within rounded brackets.

Let’s create a for loop which iterates over the expression:

for cookie in cookie_generator:
	print(cookie)

Our code returns:

Raspberry

Choc-Chip

Cinnamon

Oat

The response of this generator is the same as the one from our first example. Our syntax is significantly clearer.

It’s important to note that when you iterate over a generator that was declared using a generator expression, you don’t call it as a function. Our cookie_generator generator does not accept any input values: it already contains the code it needs to iterate over the cookie_list list.

You should only use the generator expressions syntax when you need to write a generator that performs a simple function. Printing a list of values is a good example; multiplying values in a list by a specific number is another good example.

This is because while the generator expression syntax is clear, it’s mainly designed for one-line expressions. If your generator will use more than one line of code, write it as a generator function using the syntax we discussed earlier in this article.

Why Are Generators Used?

The primary reason generators are used is that they are more concise than iterators.

Generators are like a function and do not need any __init__, __iter__, or __next__ statements to work. This is unlike an iterator which requires all three of those statements. This behavior means that you can implement a generator in significantly fewer lines of code than you could if you were writing an iterator.

What’s more, generators are efficient at using memory. Generators only produce one item at a time. They don’t create an entire sequence in memory before returning a result. This makes generators very practical if you need to iterate over a large list of data.

The way this works is that iterators use a lazy evaluation method. They only generate the next element of an iterable object when that item is requested.

Conclusion (and Challenge)

Generators allow you to create a more Pythonic iterator.

Generators make it easy to write objects using the iterator protocol without having to write __init__, __iter__, or __next__ statements. Generators are defined as functions. They use the yield statement to stop a generator and pass a value back to the main program before returning to its regular operations.

If you’re looking for a challenge, here are a few ideas:

• Write a generator that reverses a string.
• Write a generator that only returns values that contain a specific word from a list of values.
• Write a generator that multiplies every number in a list by two.

For further reading, check out our tutorial on how to write and use Python iterators. Now you’re ready to start working with generators in Python like an expert!