Generators | Python

Introduction

Generators are used to create iterators. Simply speaking, a generator is a function that returns an object (iterator) which we can iterate over (one value at a time). If a function contains at least one yield statement (it may contain other yield or return statements), it becomes a generator function. Both yield and return will return some value from a function.

Generator function vs Normal function

Difference between generator function and normal function.

Generator function contains one or more yield statement.
When called, it returns an object (iterator) but does not start execution immediately.
Methods like __iter__() and __next__() are implemented automatically. So we can iterate through the items using next().
Once the function yields, the function is paused and the control is transferred to the caller.
Local variables and their states are remembered between successive calls.

Example

# A simple generator function
def my_gen():
  n = 1
  print('This is printed first')
  # Generator function contains yield statements
  yield n

  n += 1
  print('This is printed second')
  yield n

  n += 1
  print('This is printed at last')
  yield n
  
  
  
# Returns an object but does not start execution immediately.
a = my_gen()

# We can iterate through the items using next().
next(a)

# Once the function yields, the function is paused and the control is transferred to the caller.
# Local variables and theirs states are remembered between successive calls.
next(a)
next(a)

# Finally, when the function terminates, StopIteration is raised automatically on further calls.
next(a)

# Output
This is printed first
This is printed second
This is printed at last
Traceback (most recent call last):
    next(a)
StopIteration

my_gen() looks like a typical function definition, except for the Python yield statement and the code that follows it. yield indicates where a value is sent back to the caller, but unlike return, you don’t exit the function afterward.

Generators with Loop

Generator functions are implemented with a loop having a suitable terminating condition.

def rev_str(my_str):
  length = len(my_str)
  for i in range(length - 1,-1,-1):
    yield my_str[i]

for char in rev_str("hello"):
  print(char)
  
  
# Output
o
l
l
e
h

Generator Expression

Simple generators can be created on the fly using generator expressions. Same as lambda function creates an anonymous function, generator expression creates an anonymous generator function. They are kind of lazy, producing items only when asked for. For this reason, a generator expression is much more memory efficient than an equivalent list comprehension.

# Initialize the list
my_list = [1, 3, 6, 10]

# square each term using list comprehension
a = [x**2 for x in my_list]

# same thing can be done using generator expression
a = (x**2 for x in my_list)

# Output: 1
print(next(a))

# Output: 9
print(next(a))

# Output: 36
print(next(a))

# Output: 100
print(next(a))

# Output: StopIteration
next(a)