capture changes made to object outside yield statement python

Question

I need a pythonic way to allow for this to work

def generator():
    i=0
    yield(i)
    print(i)

for i in generator():
    i=2
##should print 2

How can I modify the provided Python code to ensure that it prints 2 instead of 1 when iterating over the generator? I need a more pythonic approach for this requirement, as I have specific constraints related to ensuring the object is placed in the correct database based on changed values that are both efficient (as my database stores related objects near each other) and difficult to modify from outside.

Answer 1

You can create an iterable class to allow better control over the state of the iterator through additional methods:

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

    def set_value(self, i):
        self.i = i

    def __iter__(self):
        return self

    def __next__(self):
        value = self.i
        self.i += 1
        return value

iterator = SettableIterator()
for i in iterator:
    print(i)
    if i == 0:
        iterator.set_value(2)
    elif i == 3:
        break

This outputs:

0
2
3