Как сбросить переменную в бесконечном генераторе?

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

Answer 2

Способ с send, который позволяет (пере) установить i на желаемое значение:

def index_generator():
    i =  0
    while True:
        reset = yield i**2
        if reset is None:
            i += 1
        else:
            i = reset - 1

Использование демонстрации:

my_gen = index_generator()
print([next(my_gen) for _ in range(5)])

my_gen.send(0)
print([next(my_gen) for _ in range(5)])

my_gen.send(100)
print([next(my_gen) for _ in range(5)])

Выход:

[0, 1, 4, 9, 16]
[0, 1, 4, 9, 16]
[10000, 10201, 10404, 10609, 10816]

Вариант, который не использует трюк - 1, а вместо этого просто дает:

def index_generator():
    i =  0
    while True:
        reset = yield i**2
        if reset is None:
            i += 1
        else:
            i = reset
            yield

Answer 3

Можете ли вы сделать что-то подобное вместо этого?

for i, val in enumerate(index_generator()):
     print(val)
     if i > 1000:
             break

для каждого раза, когда вам нужно его использовать?

for для l oop неявно вызывает следующий метод, так что это немного короче, и генератор «сбрасывается» каждый раз.