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