CODEWITHSUNDEEP
pythongenerator
TheStrangeQuark
TheStrangeQuark

Reputation: 2405

iterate infinite generator until condition is met

I'm having a little trouble understanding generators in python. Lets say I have the following simple generator:

def test_gen():
    yield 1
    yield 2
    yield 5
    yield 8
    yield 12

That will just yield some numbers. Lets say I want to keep printing values from the generator until the value is more than 10. My first attempt was just

gen = test_gen()
for i in gen:
    while i > 10:
        print(i)

But this just prints 1 forever. I found that doing

gen = test_gen()
for i in gen:
    if i > 10:
        print(i)

works as expected because it just iterates through each value in the generator until StopIteration is raised. I then ran into the problem of an infinite generator such as this one that generates prime numbers:

def prime_gen():
    D = dict()
    n = 2
    while True:
        if n not in D:
            yield n
            D[n*n] = [n]
        else:
            for p in D[n]:
                D.setdefault(p + n, []).append(p)
            del D[n]
        n += 1

Then if I do something like

primes = prime_gen()
for p in primes:
    if p < 100:
        print(p)

then it will print everything up to p, but then get hung up. I think it's getting hung up because it's trying to check every value that primes generates, which going on forever. What is the correct way to iterate through an infinite generator until a condition is met like this? The only way I have found is doing

primes = prime_gen()
for p in primes:
    if p < 100:
        print(p)
    else:
        break

but I feel like there's a more pythonic way of doing it.

Upvotes: 6

Views: 2090

Answers (2)

Ignacio Vazquez-Abrams
Ignacio Vazquez-Abrams

Reputation: 798706

itertools.takewhile()

for i in itertools.takewhile(lambda x: x < 100, primes):
  print(i)

Upvotes: 14

Joe Iddon
Joe Iddon

Reputation: 20424

Your break method is quite Pythonic, but I would consider wrapping the generator:

for p in (i if i < 100 else StopIteration for i in primes):
    print(p)

Upvotes: 0

Related Questions