CODEWITHSUNDEEP
pythongenerator
mr.robot
mr.robot

Reputation: 31

How to reset a variable in an infinite generator?

I have an infinite generator object which yields some numbers, like below:

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, ...

Now I need to reset variable i to zero after some iterations, to use the generator again; let's say 100 more times in my code. I can't repeat my_gen = index_generator() every time. So, any solutions?

Upvotes: 3

Views: 428

Answers (3)

Kelly Bundy
Kelly Bundy

Reputation: 27609

A way with send which allows to (re)set i to a desired value:

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

Demo usage:

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)])

Output:

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

A variation that doesn't use the - 1 trick but instead just yields:

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

Upvotes: 2

kaya3
kaya3

Reputation: 51063

One option is to write a class with a __next__ method, so it can be used as an iterator, and a reset method which resets i. The __iter__ method returning self is also needed if you aren't just calling next manually.

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

Usage example:

>>> 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

Upvotes: 6

OlympusMonds
OlympusMonds

Reputation: 158

Can you do something like this instead?

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

for each time you need to use it?

The for loop calls the next method implicitly, so this is a bit shorter, and the generator 'resets' each time.

Upvotes: 1

Related Questions