Infinite prime generator

This code acts as an infinite generator of prime numbers.

As new prime numbers are found, they are added to a set. Then, a number x is found to be prime if none of the numbers in the prime set are factors of x. Only the primes less than or equal to the square root of x need to be checked.

import itertools
from math import sqrt

class stream:
""" Class of infinite streams. """

def prime():
""" Stream of prime numbers. """
prime_set = {2} # Set of prime numbers that have been found
yield 2 # First prime
for x in itertools.count(3, 2): # Check odd numbers, starting with 3
primes_below_sqrt = {i for i in prime_set if i <= sqrt(x)}
for prime in primes_below_sqrt:
if x % prime == 0:
break # x is divisible by a prime factor, so it is not prime
else:
prime_set.add(x) # x has been shown to be prime
yield x

Using the itertools recipe:

def take(iterable, n):
""" Returns first n items of the iterable as a list. """
return list(itertools.islice(iterable, n))

Output:

>>> take(stream.prime(), 10)
[2, 3, 5, 7, 11, 13, 17, 19, 23, 29]

Logic mostly looks sound. Some general Python style comments:

• I see no reason to put this function on a class – just make it a top-level function and be done with it. Even if you’re writing multiple infinite streams, group them in a single file, not on a class.

• Read PEP 8, the Python style guide. In particular, two spaces before inline comments, class names are uppercase.

• Your comments aren’t very helpful; they really just describe what the code is doing. It would be better to explain why the code was written that way. For example, replace:

# Check odd numbers, starting with 3

with:

# We only need to check odd numbers, because all evens > 2 are not prime

• Use keyword arguments for clarity. Your call to itertools.count() is clearer like so:

itertools.count(start=3, step=2)

I was able to sort-of guess based on the comment, but keyword arguments always improve clarity and reduce ambiguity.

This is a performance issue:

primes_below_sqrt = {i for i in prime_set if i <= sqrt(x)}

By using a set comprehension instead of a generator expression, you are actually making a copy of the subset.

• Would the problem be fixed just by switching the curly brackets to round brackets? – Vermillion Nov 3 '16 at 20:25
• Yes, that could fix the issue. – 200_success Nov 3 '16 at 20:27