python迭代器(Iterator)与生成器(Generator)

最近在阅读技术文档时遇到迭代器和生成器相关的概念，这篇文章总结一下学到的东西。

感谢：

1. Python进阶——什么是迭代器？

2. Python __iter__()迭代器 （十三）_鼻涕虫de皮皮的博客-CSDN博客

3. python 类的__iter__、__next__方法与yield_JpyJpy321的博客-CSDN博客

主要结合实例，以个人理解总结"容器"，"迭代器"，"可迭代对象"，"生成器"这四个概念及其用法。

1. 容器，在Python中主要指用类实现的一种数据结构，这个类的定义一定是包含了__contain__()方法，以支持Python的in关键字能够和这个类兼容，确定某个元素是否在容器里，如：

class CustomContainer:
    def __init__(self) -> None:
        self.data = [1,2,3,4,5]
    
    def __contains__(self, element):
        if element in self.data:
            return True
        else:
            return False


if __name__ == "__main__":
    # test container:
    container = CustomContainer()
    print("1 in container: ")
    print(1 in container)

    print("10 in container: ")
    print(10 in container)
    

2. 迭代器，在Python指用类实现的服从迭代器协议的数据结构，即迭代器类定义了__iter__()方法和__next__()方法，使得该迭代器能够被for关键字迭代等等，其中：

__iter__()方法返回一个迭代器对象，一般是self，这个返回的对象必须实现了__next__()方法，因为for循环使用这个迭代器时，会首先调用一次__iter__()，然后每次循环时调用这个__iter__()返回的对象的__next__()方法，将__next__()返回值赋给循环变量。当数据结构中的数据被遍历过一遍时，要__next__()负责抛出StopIteration异常，告诉for语句停止迭代。

另外，关于一个迭代器对象可以迭代几次的问题。这个取决于抛出StopIteration异常前，是否有将这个迭代器对象的状态重置。如果没有重置，第一次for循环不会出现问题，第二次不会报错，但是for语句块没有循环。下面将通过实际例子说明

循环后无重置的Iterator:

class CustomIterator:
    def __init__(self) -> None:
        self.data = [1,2,3,4,5]
        self.current_idx = 0
    
    def __iter__(self):
        return self
    
    def __next__(self):
        if self.current_idx < len(self.data):
            self.current_idx += 1
            return self.data[self.current_idx - 1]
        else:
            raise StopIteration()

if __name__ == "__main__":
    # test Iterator:
    iterator = CustomIterator()
    print('first loop: ')
    for idx,element in enumerate(iterator):
        print('{}-th element'.format(idx + 1))
        print(element)
    print('first loop end')
    print('second loop: ')
    for idx,element in enumerate(iterator):
        print('{}-th element'.format(idx + 1))
        print(element)
    print('second loop end')

循环后有重置的Iterator(在__next__()的raise StopIteration语句前进行处理):

class CustomIterator:
    def __init__(self) -> None:
        self.data = [1,2,3,4,5]
        self.current_idx = 0
    
    def __iter__(self):
        return self
    
    def __next__(self):
        if self.current_idx < len(self.data):
            self.current_idx += 1
            return self.data[self.current_idx - 1]
        else:
            self.current_idx = 0 # 对Iterator的状态进行重置
            raise StopIteration()

if __name__ == "__main__":
    # test Iterator:
    iterator = CustomIterator()
    print('first loop: ')
    for idx,element in enumerate(iterator):
        print('{}-th element'.format(idx + 1))
        print(element)
    print('first loop end')
    print('second loop: ')
    for idx,element in enumerate(iterator):
        print('{}-th element'.format(idx + 1))
        print(element)
    print('second loop end')

3. 可迭代对象，指__iter__()可以返回一个迭代器的对象，常用的list，tuple，dict都是可迭代对象，这种对象可以将__next__()的实现交给其它类来处理，如这个WeirdIterator类：

class CustomIterator:
    def __init__(self) -> None:
        self.data = [1,2,3,4,5]
        self.current_idx = 0
    
    def __iter__(self):
        return self
    
    def __next__(self):
        if self.current_idx < len(self.data):
            self.current_idx += 1
            return self.data[self.current_idx - 1]
        else:
            self.current_idx = 0 # 对Iterator的状态进行重置
            raise StopIteration()

class WeirdIterator:
    def __init__(self) -> None:
        pass
    
    def __iter__(self):
        return CustomIterator()

if __name__ == "__main__":
    # test WeirdIterator:
    iterator = WeirdIterator()
    print('loop: ')
    for idx,element in enumerate(iterator):
        print('{}-th element'.format(idx + 1))
        print(element)
    print('loop end')

 4. 生成器，生成器是一种特殊的迭代器，可以通过生成器表达式或生成器函数实现，每次循环会保存生成器当前的状态，下次循环时当前生成器中的状态推导下一个元素

生成器表达式指用用圆括号括住推导式，这个表达式可以返回一个迭代器对象，如：

if __name__ == "__main__":
    # test Generator:
    generator_1 = (i for i in range(1,6))
    print('loop: ')
    for idx, element in enumerate(generator_1):
        print('{}-th element'.format(idx + 1))
        print(element)
    print('loop end')

生成器函数采用yield关键字替代return关键字。需要注意的是，生成器函数这个名称有误导作用。实际上，替换以后，函数名实际是个类名，函数名加括号返回的是一个生成器对象，如：

def funcGenerator(lst:list):
    for i in lst:
        yield i

if __name__ == "__main__":
    # test Generator:
    generator_2 = funcGenerator([1,2,3,4,5])
    print('loop: ')
    for idx, element in enumerate(generator_2):
        print('{}-th element'.format(idx + 1))
        print(element)
    print('loop end')

以上两种生成器的写法产生的生成器对象也有定义__iter__()以及__next__()方法。

但是以上两种写法也存在问题，即只能被for循环遍历一次，由于状态没有重置，第二次遍历会失败。一种解决方案是每次for循环都使用一个新的生成器对象，借助前面介绍的可迭代对象的特性，可以这样写：

class CustomGenerator:
    def __init__(self) -> None:
        self.curr_data = 0

    def __iter__(self):
        while self.curr_data < 5:
            self.curr_data += 1
            yield self.curr_data
        
        self.curr_data = 0

if __name__ == "__main__":

    # test Generator:
    generator_3 = CustomGenerator()
    print("loop 1: ")
    for idx, element in enumerate(generator_3):
        print('{}-th element'.format(idx + 1))
        print(element)
    print('loop 1 end')

    print("loop 2: ")
    for idx, element in enumerate(generator_3):
        print('{}-th element'.format(idx + 1))
        print(element)
    print('loop 2 end')

如果数据比较多，一般的迭代器将所有数据读入内存，会导致内存占用过多。而生成器主要可以用于解决这个问题