Title: Use the yield statement to generate primes in Python
Python's yield statement creates a generator function that can return a value and then suspend execution until the next time a value is required. The program can iterate through the results and the generator only produces values each time it is asked for one.
The get_primes function shown in the following code uses yield to return prime numbers.
def get_primes():
yield 2
odd_primes = []
number = 1
while True:
# Move to the next odd number.
number += 2
# See if number is divisible by any of the primes.
is_prime = True
for p in odd_primes:
if number % p == 0:
is_prime = False
break
if p * p >= number:
break
if is_prime:
odd_primes.append(number)
yield number
The function first yields the value 2 because that is the first prime number. It then creates the list odd_primes to hold the primes that it will generate. It sets variable number to 1 and then enters an infinite loop.
Each time through the loop, the code adds 2 to number so it holds the next odd number. The function then enters a loop where it sees whether any of the previously discovered primes in the odd_primes list divides evenly into number. If any prime divides into number, then number is not prime so the program sets is_prime to False.
If the test prime p squared is larger than number, then the code breaks out of the inner loop. (For example, there's no point testing whether 13 divides evenly into 100 because 13 * 13 > 100.)
After the inner loop ends, if is_prime is still True, we know that number is prime. In that case, the function adds it to the odd_primes list and yields it.
You'll notice that the function never leaves its infinite loop so you may wonder how control returns to the calling code. The answer is the yield statement. Each time the program reaches a yield statement, the generator function returns a value and then the function stops until the next time it is invoked. It then resumes execution where it left off until it reaches another yield statement. The process continues indefinitely until the calling code stops calling the generator or the generator exits.
The following code shows the example's test program.
for p in get_primes():
if p > 100: break
print(p, end=' ')
print()
for p in get_primes():
print(p)
another = input('Another (y/n)? ')
if another == 'n': break
This code first uses get_primes and then loops variable p through the results. If p is greater than 100, the code breaks out of the loop. Otherwise it prints p and continues the loop. The result is this part of the code displays the prime numbers between 0 and 100.
To show that the function pauses after each yield statement, the code calls get_primes again and loops through the results. This time it displays each number generated and then pauses for input. If you type n and press Enter, the program breaks out of its loop and stops. This shows that the generator isn't generating the next value until it is needed.
You can also place print statements inside get_primes to see when it executes each statement. You'll find that it pauses after each yield statement and waits until the next value is required.
Also note that calling get_primes a second time creates a new generator that starts from scratch. In fact, you can even nest loops through multiple generators like this.
for p in get_primes():
if p > 20: break
print(f'*** {p} ***')
for q in get_primes():
if q > 10: break
print(f'{q}')
Download the example to see additional details.
|