CODEWITHSUNDEEP
pythonpython-3.x
Kyle G
Kyle G

Reputation: 1038

Setting a limit on an infinite generator in Python

I'm trying to make the below generator to be able to set an upper limit on the numbers returned.

Calling list(it.takewhile(lambda x: x < 100, get_primes())) returns a list of all primes under 100 as expected but list(get_primes(100)) (which should return the same list in the same way) just returns with an empty list.

Obviously, I could include an if n and candidate>=n: break in the for loop but I'm mostly interested in why the if n: return construct doesn't work like I'm expecting it should. Shouldn't it just return the same takewhile iterator that works above? What am I overlooking here?

import itertools as it

def get_primes(n=None):
    """
    Generates primes to a max of n.

    >>> list(get_primes(100))
    [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97]
    """
    if n:
        return it.takewhile(lambda x: x < n, get_primes())
    composites = {}
    yield 2
    for candidate in it.count(3, 2):
        prime_factor = composites.pop(candidate, None)
        if prime_factor is None:
            yield candidate
            composites[candidate**2] = candidate
        else:
            composite = candidate + 2*prime_factor
            while composite in composites:
                composite += 2*prime_factor
            composites[composite] = prime_factor

Upvotes: 5

Views: 3532

Answers (2)

Roukanken
Roukanken

Reputation: 116

The problem is with the line 

return it.takewhile(lambda x: x < n, get_primes())

Since it's a generator, returning something stops the execution, and raises StopIteration(). You need to return the values in generator

#return all the values from generator
for a in it.takewhile(lambda x: x < n, get_primes())
    yield a

return

Upvotes: 2

NPE
NPE

Reputation: 500437

Here:

return it.takewhile(lambda x: x < n, get_primes())

Since this is a generator, it needs to yield these values instead of returning them. Depending on your Python version, you might be able to use the yield from syntax.

The following might be useful as background reading: Return in generator together with yield in Python 3.3

Upvotes: 3

Related Questions