如何在无限生成器中重置变量？

Question

我有一个无限生成器对象，它产生一些数字，如下所示：

def index_generator():
    i = 0
    while True:
        yield i**2
        i += 1

my_gen = index_generator()

for i in range(100):
    print(next(my_gen))  # outputs 0,1,4,9, ...

现在我需要i在一些迭代后将变量重置为零，以再次使用生成器；让我们在我的代码中再说 100 次。我不能my_gen = index_generator()每次都重复。那么，有什么解决办法吗？

Answer 1

一种选择是用一个__next__方法编写一个类，这样它就可以用作迭代器，以及一个reset重置i. 该__iter__方法返回self也需要，如果你不只是打电话next手动。

class IndexGenerator:
    def __init__(self):
        self.i = 0

    def __next__(self):
        r = self.i ** 2
        self.i += 1
        return r

    def reset(self):
        self.i = 0

    def __iter__(self):
        return self

用法示例：

>>> g = IndexGenerator()
>>> next(g)
0
>>> next(g)
1
>>> next(g)
4
>>> next(g)
9
>>> g.reset()
>>> next(g)
0
>>> next(g)
1
>>> for x in g:
...     print(x)
...     if x > 10: break
... 
4
9
16