python - 迭代无限生成器直到满足条件

Question

我在理解 python 中的生成器时遇到了一些困难。假设我有以下简单的生成器:

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

这只会产生一些数字。假设我想一直打印生成器的值，直到值超过 10。我的第一次尝试只是

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

但这只会永远打印 1。我发现这样做

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

按预期工作，因为它只是遍历生成器中的每个值，直到引发 StopIteration。然后我遇到了无限生成器的问题，例如这个生成素数的生成器:

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

然后如果我做类似的事情

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

然后它将打印直到 p 的所有内容，但随后挂断了。我认为它挂断了，因为它试图检查 primes 生成的每个值，这将永远持续下去。迭代无限生成器直到满足这样的条件的正确方法是什么？我发现的唯一方法是做

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

但我觉得有一种更 pythonic 的方式来做这件事。

Answer 1

itertools.takewhile()

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