python 迭代器 iter多次消费

问题

Python 中的迭代器是我们经常使用的迭代工具， 但其只能消费一次，再次消费便会出现 StopIteration 报错。

解决方案

封装了一个类，当迭代器使用完后再次初始化。

代码

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)