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
                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))


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

2 Answers 2


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


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

  • \$\begingroup\$ Would the problem be fixed just by switching the curly brackets to round brackets? \$\endgroup\$
    – Vermillion
    Nov 3, 2016 at 20:25
  • 1
    \$\begingroup\$ Yes, that could fix the issue. \$\endgroup\$ Nov 3, 2016 at 20:27

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.