Python Generators

A generator is a function that can pause and later carry on.

Calling a generator function gives you back a generator object, and this object is an iterator.

None of its code has run yet. It only starts when you iterate over the generator.

Example

A basic generator function:

Write a Python program that defines a generator function my_generator() which yields 1, 2, and 3.
Use a for loop to iterate over the generator and print each value.

> Output: complete Python program

Prompt

Write a Python program that defines a generator function my_generator() which yields 1, 2, and 3. Use a for loop to iterate over the generator and print each value.

def my_generator():
  yield 1
  yield 2
  yield 3

for value in my_generator():
  print(value)

The yield Keyword

What turns a function into a generator is the yield keyword.

On reaching a yield, the function hands back a value and saves its state. When the generator is asked for another value, it picks up right where it stopped.

return ends a function for good, but yield only pauses it, so a function can use it many times.

Example

A generator yielding a series of numbers:

Write a Python program that defines a generator count_up_to(n) which yields numbers from 1 to n using a while loop.
Use a for loop to print each number up to 5.

> Output: complete Python program

Prompt

Write a Python program that defines a generator count_up_to(n) which yields numbers from 1 to n using a while loop. Use a for loop to print each number up to 5.

def count_up_to(n):
  count = 1
  while count <= n:
    yield count
    count += 1

for num in count_up_to(5):
  print(num)

Generators Save Memory

A generator produces values as needed instead of keeping them all in memory.

That makes a real difference with large datasets:

Example

A generator for a very long sequence:

Write a Python program that defines a generator large_sequence(n) which yields numbers from 0 to n-1.
Create a generator for 1000000 numbers and print the first 3 values using next().

> Output: complete Python program

Prompt

Write a Python program that defines a generator large_sequence(n) which yields numbers from 0 to n-1. Create a generator for 1000000 numbers and print the first 3 values using next().

def large_sequence(n):
  for i in range(n):
    yield i

# This doesn't create a million numbers in memory
gen = large_sequence(1000000)
print(next(gen))
print(next(gen))
print(next(gen))

Using next() with Generators

The next() function lets you step through a generator by hand.

Once all the values have been yielded, another call raises a StopIteration exception.

Example

Write a Python program that defines a generator simple_gen() which yields "Emil", "Tobias", and "Linus".
Use next() three times to print each value.

> Output: complete Python program

Prompt

Write a Python program that defines a generator simple_gen() which yields "Emil", "Tobias", and "Linus". Use next() three times to print each value.

def simple_gen():
  yield "Emil"
  yield "Tobias"
  yield "Linus"

gen = simple_gen()
print(next(gen))
print(next(gen))
print(next(gen))

Example

Write a Python program that defines a generator simple_gen() which yields 1 and 2.
Use next() three times to demonstrate the StopIteration exception on the third call.

> Output: complete Python program

Prompt

Write a Python program that defines a generator simple_gen() which yields 1 and 2. Use next() three times to demonstrate the StopIteration exception on the third call.

def simple_gen():
  yield 1
  yield 2

gen = simple_gen()
print(next(gen))
print(next(gen))
print(next(gen)) # This will raise StopIteration

Generator Expressions

Generator expressions create generators with syntax close to list comprehensions. The only difference is parentheses in place of square brackets:

Example

A list comprehension next to a generator expression:

Write a Python program that creates a list comprehension [x * x for x in range(5)] and prints it.
Then create a generator expression (x * x for x in range(5)), print the generator object, and print its values as a list.

> Output: complete Python program

Prompt

Write a Python program that creates a list comprehension [x * x for x in range(5)] and prints it. Then create a generator expression (x * x for x in range(5)), print the generator object, and print its values as a list.

# List comprehension - creates a list
list_comp = [x * x for x in range(5)]
print(list_comp)

# Generator expression - creates a generator
gen_exp = (x * x for x in range(5))
print(gen_exp)
print(list(gen_exp))

Example

Passing a generator expression to sum:

Write a Python program that calculates the sum of squares from 0 to 9 using sum() with a generator expression.
Print the total.

> Output: complete Python program

Prompt

Write a Python program that calculates the sum of squares from 0 to 9 using sum() with a generator expression. Print the total.

# Calculate sum of squares without creating a list
total = sum(x * x for x in range(10))
print(total)

Fibonacci Sequence Generator

The Fibonacci sequence is a natural fit for a generator.

It can go on producing values indefinitely without exhausting your memory:

Example

Print the first 100 Fibonacci numbers:

Write a Python program that defines an infinite Fibonacci generator using yield.
Print the first 100 Fibonacci numbers using next().

> Output: complete Python program

Prompt

Write a Python program that defines an infinite Fibonacci generator using yield. Print the first 100 Fibonacci numbers using next().

def fibonacci():
  a, b = 0, 1
  while True:
    yield a
    a, b = b, a + b

# Get first 100 Fibonacci numbers
gen = fibonacci()
for _ in range(100):
  print(next(gen))

Generator Methods

For more advanced control, generators come with some special methods:

Example

With the send() method, you can pass a value into the generator:

Write a Python program that defines an echo_generator() which receives values via yield and prints them.
Prime the generator with next(), then send "Hello" and "World".

> Output: complete Python program

Prompt

Write a Python program that defines an echo_generator() which receives values via yield and prints them. Prime the generator with next(), then send "Hello" and "World".

def echo_generator():
  while True:
    received = yield
    print("Received:", received)

gen = echo_generator()
next(gen) # Prime the generator
gen.send("Hello")
gen.send("World")

Example

The close() method shuts the generator down:

Write a Python program that defines a generator my_gen() with a try/finally block that yields 1, 2, 3 and prints "Generator closed" in finally.
Get the first value with next(), then close the generator.

> Output: complete Python program

Prompt

Write a Python program that defines a generator my_gen() with a try/finally block that yields 1, 2, 3 and prints "Generator closed" in finally. Get the first value with next(), then close the generator.

def my_gen():
  try:
    yield 1
    yield 2
    yield 3
  finally:
    print("Generator closed")

gen = my_gen()
print(next(gen))
gen.close()