Computer Science Wiki

Generator and the yield Keyword

4 min read

The yield keyword is used to create generators, which are special types of iterators that allow values to be produced lazily, one at a time, instead of storing them all in the memory and return them all at once. The yield keyword is more memory efficient since it uses lazy iteration, where all values are not computed at once, but generate one value at a time on demand. This is more preferred than using the return keyword in scenarios when you are processing large data streams, reading big files line by line, or any time you want efficient memory usage.

How yield works

  • When a function contains yield, it becomes a generator function
  • Calling the function does not execute it immediately. Instead, it returns (yields) a generator object
  • The next() method on the generator resumes execution from where yield last returned a value
  • The generator remembers its state between calls

The yield keyword temporarily freezes the execution and returns the value to the caller. Because it freezes the execution, the function’s state is retained, allowing it continue from where it left off; unlike the return keyword which stops further execution of the function.

def generate_numbers():
	num = 0
	while True:
		yield num
		num += 1
 
number = generate_numbers()
print(type(number))
<class 'generator'>

From the above example, we know that the generate_number() function returns a generator, since it contains the yield keyword. Now we try to print the numbers using the generator function.

def generate_numbers():
	num = 0
	while True:
		yield num
		num += 1
 
number = generate_numbers()
print(next(number)) # output: 0
print(next(number)) # output: 1
print(next(number)) # output: 2

Let’s inspect its underlying behaviour. When the first next() function is called, the execution runs until it reaches the yield statement after which the number is generated and returned to the caller, and the execution flow is suspended. On the second iteration (the second call to the next() function), the execution continued from where it left off, incrementing the number by one (num += 1), and continue the while loop.

def generate_numbers():
	num = 0
	while True:
		yield num # Program runs until it encounters the yield statement
		# The line below is suspended
		num += 1

Generator object and the iterator protocol

The generator object implements the iterator protocol:

  • __iter__(): returns the iterator object itself
  • __next__(): returns the next value or raises StopIteration when exhausted

You don’t need to manually implement __iter__() and __next__(), Python does it automatically. When you call a generator function, it does not execute immediately. Instead:

  1. It creates a generator object
  2. The generator implements the iterator protocol
  3. When you call the next(generator), execution resumes from where it was paused at yeild

When using a for loop or a direct call to the next()function, Python automatically calls the __next__() method of the generator object.

for i in range(5):
	yield i # This pauses execution and remembers state
0
1
2
3
4

def count_up_to(n):
    count = 1
    while count <= n:
        yield count  # Pauses and returns the value
        count += 1   # Resumes from here when next() is called
 
gen = count_up_to(3)
print(next(gen))  # Output: 1
print(next(gen))  # Output: 2
print(next(gen))  # Output: 3
# print(next(gen))  # Raises StopIteration

Manually implementing a generator

Since a generator automatically creates an iterator, let’s mimic the behaviour manually:

class MyGenerator:
	def __init__(self, n):
		self.n = n
		self.current = 0 # Track current state
		
	def __iter__(self):
		return self # An iterator returns itself
 
	def __next__(self):
		if self.current < self.n:
			value = self.current
			self.current += 1
			return value # Acts like yield
		else:
			raise StopIteration
 
# Client
gen = MyGenerator(3)
for num in gen:
	print(num)
0
1
2

Example - without yield

The below function uses a list to pre-store the values to be printed.

def fun(m):
	my_list = []
	for i in range(m):
		my_list.append(i)
	return my_list
 
for num in fun(5):
	print(num)
0
1
2
3
4

Example - with yield keyword

The fun(m) function is a generator function (since it contains the yield keyword) generates the numbers one at a time instead of storing them all in the memory.

def fun(m):
	for i in range(m):
		yield i  # yeild value one by one
 
for num in fun(5):
	print(num)
0
1
2
3
4

Back to parent page: Web and Application Development

Reference:

Graph View

Backlinks