Setting a limit on an infinite generator in Python

Question

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

score 3 · Accepted Answer · edited May 23 '17 at 12:27

3

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

edited May 23 '17 at 12:27

Community

1
1

answered Oct 05 '14 at 19:40

NPE

486,780
108
951
1,012

1

though replacing `return` with `yield` here won't have the expected behaviour (the generator will yield another generator as first value). – isedev Oct 05 '14 at 19:44
@isedev: `yield from` will be a straight replacement for `return` (a plain `yield`, of course, will not). – NPE Oct 05 '14 at 19:46
Sure, agreed if you can use `yield from`. Otherwise, the actual generator will need to be wrapped in a nested function so the outer function can `return` in both cases. – isedev Oct 05 '14 at 19:46
Makes sense. I guess I had assumed that because the `return` statement would happen before any `yield` statement, it would happily return my takewhile iterator. From the linked thread, it appears that it doesn't work that way. Thanks. – Kyle G Oct 05 '14 at 20:05

score 2 · Answer 2 · answered Oct 05 '14 at 19:48

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

Setting a limit on an infinite generator in Python

2 Answers2