Memoized Prime Generator

Question

Here is my code:

def odd_numbers_from(n):
    if n % 2 == 0:
        n = n + 1
    while True:
        yield n
        n = n + 2

def takeWhile(pred, lst):

    for item in lst:
        if pred(item):
            yield item
        else:
            break

def takePrimes(iterator):

    for n in iterator:

        if not 0 in (n % k for k in takeWhile(lambda x: x**2 <= n, primes())):

            yield n

# -- memoized prime generator
def primes(earlier_primes=[2]):

    yield from earlier_primes

    for new_prime in takePrimes(odd_numbers_from(earlier_primes[-1]+1)):

        earlier_primes.append(new_prime)

        yield new_prime

if __name__ == "__main__":

    for k in primes():

        print(k)
        if k > 300:
            break

In first run, it will produce prime numbers as long as they are requested. On subsequent runs, it will just yield from previously generated primes. We can test this behaviour like this:

# -- memoized prime generator
def primes(earlier_primes=[2]):

    yield from earlier_primes

    for new_prime in takePrimes(odd_numbers_from(earlier_primes[-1]+1)):
        print("generated new prime")
        earlier_primes.append(new_prime)

        yield new_prime

if __name__ == "__main__":

    for k in primes():

        print(k)
        if k > 20:
            break

    print("=====")

    for k in primes():

        print(k)
        if k > 30:
            break

This should output:

2
generated new prime
3
generated new prime
5
generated new prime
7
generated new prime
11
generated new prime
13
generated new prime
17
generated new prime
19
generated new prime
23
=====
2
3
5
7
11
13
17
19
23
generated new prime
29
generated new prime
31

How does it look?

Answer 1

1. Review

1. There are no docstrings. What do these functions do? How do I call them? What do they return?

2. The function odd_numbers_from could be implemented using itertools.count, like this:

   def odd_numbers_from(n):
       """Generate the odd numbers greater than or equal to n."""
       return itertools.count(n + 1 - (n % 2), 2)
   

3. But if you started with earlier_primes=[2,3] then you could avoid the special case here.

4. The function takeWhile is built into Python under the name itertools.takewhile.

5. takePrimes is only called from primes, so would make more sense to be inlined there.

6. A Python object that can be looped over using the for statement (like the argument to takePrimes) is properly known as an iterable, not a iterator.

7. Instead of not 0 in expression, write all(expression).

2. Revised code

from itertools import count, islice, takewhile

def primes(earlier_primes=[2, 3]):
    """Generate the prime numbers.

    >>> list(islice(primes(), 10))
    [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]

    """
    yield from earlier_primes
    for n in count(earlier_primes[-1] + 2, 2):
        if all(n % p for p in takewhile(lambda p: p**2 <= n, primes())):
            earlier_primes.append(n)
            yield n

Answer 2

Don't reinvent the weel

In general Python already has all the general purpose functions built-in:

def takeWhile(pred, lst):
    for item in lst:
        if pred(item):
            yield item
        else:
            break

:

from itertools import takewhile, count

And it will work exactly like your implementation.

Use less blank lines

Blank lines separate logical blocks of meaning, use them sparingly, for example primes should be:

def primes(earlier_primes=[2]):
    yield from earlier_primes
    for new_prime in takePrimes(odd_numbers_from(earlier_primes[-1]+1)):
        earlier_primes.append(new_prime)
        yield new_prime

Use docstrings

Comments such as: # -- memoized prime generator belong after the function definition as a docstring, like:

def primes(earlier_primes=[2]):
    """"Memoized prime generator."""
    yield from earlier_primes
    for new_prime in takePrimes(odd_numbers_from(earlier_primes[-1]+1)):
        earlier_primes.append(new_prime)
        yield new_prime

Modularize more

Your code is divided in many small functions and this is very good, but you can do more, for example the line:

if not 0 in (n % k for k in takeWhile(lambda x: x**2 <= n, primes())):

is a primality check, but I would very much prefer:

def is_prime(n):
    return not 0 in (n % k for k in takeWhile(lambda x: x**2 <= n, primes()))

and then:

if is_prime(n):