python iterator iter repeatedly consumption

problem

Iterator in Python is an iterative tool we often use, but only once the consumer, the consumer will again appear StopIteration error.

solution

Encapsulates a class, when an iterator after re-initialized.

Code

class RepeatDataLoader():
    def __init__(self):
        self.data_iter = self.data_loader()
        self.renew_epoch = False

    def __next__(self):
        try:
            batch = self.data_iter.__next__()
            self.renew_epoch = False
        except StopIteration:
            self.data_iter = self.data_loader()
            batch = self.data_iter.__next__()
            self.renew_epoch = True
        return batch
    
    @staticmethod
    def data_loader():
        for i in range(5):
            yield i

def checkIter(myIter):
    for _ in range(15):
        a = next(myIter)
        print (a)

if __name__ == "__main__":
    oneIter = RepeatDataLoader.data_loader()
    repeatIter = RepeatDataLoader()
    checkIter(repeatIter)