Infinite list of prime numbers using Python generators

Sometimes we want to create collections of elements that are very expensive to calculate. The first option is to create a list and wait until all the elements are calculated before we use it. Although this works, it is not very efficient. To make it a bit more efficient, modern languages provide a way to create custom iterators so each element is calculated only when needed (this is also called lazy initialization). Also, iterators allows us to create infinite collections!

Python has, in my opinion, one of the most succint and elegant ways to declare iterators: generators.

Without further ado, let's try to create an infinite list of prime numbers.

The first thing we need is a way to detect if a number is prime:

from math import sqrt

def is_prime(n):
    if (n <= 1):
        return False
    if (n == 2):
        return True
    if (n % 2 == 0):
        return False

    i = 3
    while i <= sqrt(n):
        if n % i == 0:
            return False
        i = i + 2

    return True

Now, using our is_prime function we can do:

def prime_generator():
    n = 1
    while True:
        n += 1
        if is_prime(n):
            yield n

And that's it! Just call the function and get elements from it:

generator = prime_generator()

for i in range(10):
    print(next(generator))

Or create a list of the first N prime numbers:

from itertools import islice

array = [x for x in islice(prime_generator(), 10)]

As you could see, the iterator definition is one of the shortest and simplest among all languages.

Comments

[orenovadia]

Since the time complexity of is_prime is O(sqrt(n)):
Worth a mention that for prime numbers up to N, overall, using this method is not as time efficient as sieve methods that eagerly pre-calculate all the primes up to N (this is because sieves use previously found primes to eliminate other composite numbers, instead of performing an exhaustive search for each number).

[Rishit Chaudhary]

Hi,
Amazing post. It's a really great example to understand how to use generators in Python!

I feel the code above can be improved a bit further, by starting the check for prime numbers from a more suitable number, hence the primeGenerator() method can be rewritten as follows

def primeGenerator():
    n = 3
    while True:
        if isPrime(n):
            yield n
        n += 1

[Alejandro Bezdjian]

Hi! I’m glad you liked it! The snippet you provided should start from the number 2 because it is also a prime number.

[Rishit Chaudhary]

Yeah, your absolutely right!
I guess my previous snippet would only be useful for people wanting to print all prime numbers except the even prime number (which is 2).

Here is my corrected code, which gives all the prime numbers:

def primeGenerator():
    n = 2
    while True:
        if isPrime(n):
            yield n
        n += 1

Thank You,
Rishit

[Billy Einkamerer]

Thanks for this - it was super helpful.
I hacked it a bit and added a value limit for the prime number check, so that you can check the prime numbers up to X.

Changed up the prime generator function to have a limit (also set a default limit):

then used that limit in the call of the generator into a list:

so no longer need itertools and islice.

[Alejandro Bezdjian]

I’m glad that you find it useful 🙂

[Frank]

Nice post but 2 is not prime. You need to shift the n==2 if-statement to the top.

[Alejandro Bezdjian]

Hi Frank, I’m afraid 2 is a prime number

[Frank]

Hi Alejandro, yes, 2 is prime but your code says no.

[Bwillis]

The code returns True if it is a prime number. It says if n == 2, return True, which states that 2 is a prime number, because it is actually a prime number

[Jovan Cosovic]

Can u help me with describing the part of the code below i = 3?
Thanks.

[Alejandro Bezdjian]

Hi Jovan, the first ifs in the code cover the cases where i <= 1 and i == 2, so starting the iteration with i = 1 will repeat the first 2 cases. It's just a minimum performance improvement.