## DEV Community

Alejandro Bezdjian

Posted on • Updated on

# 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.

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

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 • Edited

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.