30天拿下Python之迭代器和生成器

概述

在上一节，我们介绍了Python的模块和包，包括：什么是模块、导入模块、自定义模块、__name__、什么是包、创建包、导入包等内容。在这一节中，我们将介绍Python的迭代器和生成器。在Python中，迭代器是一个非常重要的概念，它使得我们能够遍历一个序列而无需使用索引。迭代器不仅限于列表、元组、字符串等，我们也可以创建自定义的迭代器对象。生成器是一种特殊的迭代器，能够根据需要生成数据。与传统的列表、元组等不同，生成器可以在需要时才生成数据，从而有效节省内存空间。

使用迭代器

迭代器是Python中一个重要的设计模式。迭代器是一个能够记住遍历的位置的对象，可以让我们遍历一个容器，比如：列表、元组、字典等。Python的许多内置数据类型，比如：列表、元组、字典、集合和字符串，都实现了迭代器接口。可以使用iter()函数获取这些对象的迭代器，然后使用next()函数逐个获取元素。

info = ['hello', 'world', 'python']
# 获取迭代器
it = iter(info)
# 遍历下一个元素，输出：hello
print(next(it))
# 遍历下一个元素，输出：world
print(next(it))
# 遍历下一个元素，输出：python
print(next(it))

除了使用next()函数遍历元素外，也可以使用for语句进行遍历。

info = ['hello', 'world', 'python']
# 获取迭代器
it = iter(info)
# 使用for遍历，依次输出：hello world python
for item in it:
    print(item)

当然，也可以使用while语句结合next()函数遍历所有元素。此时，需要额外处理StopIteration异常。这是因为，当next()函数遍历完序列中的所有元素后，会抛出StopIteration异常。

info = ['hello', 'world', 'python']
# 获取迭代器
it = iter(info)
# 使用while和next函数遍历，依次输出：hello world python
while True:
    try:
        print(next(it))
    except StopIteration:
        break

创建迭代器

在Python中，可以通过定义一个包含__iter__()和__next__()函数的类来创建自定义的迭代器。其中，__iter__()函数返回一个特殊的迭代器对象，一般为迭代器对象本身；__next__()函数会返回序列中的下一个元素，并通过抛出StopIteration异常标识整个迭代过程的完成。

# 自定义迭代器类
class CustomIterator:
    def __init__(self):
        self.value = 1
    def __iter__(self):
        return self
    def __next__(self):
        if self.value <= 6:
            result = self.value
            self.value += 1
            return result
        else:
            # 遍历结束时，需要抛出StopIteration异常
            raise StopIteration
       
# 创建自定义迭代器对象
cus_iter = CustomIterator()
# 遍历自定义迭代器，依次输出：1 2 3 4 5 6
for item in cus_iter:
    print(item)

在上面的示例代码中，首先定义了一个名为CustomIterator的类。这个类包含__iter__()和__next__()函数。__iter__()函数返回迭代器对象本身，而__next__()函数返回序列中的下一个元素。在__next__()函数中，检查当前值是否小于等于6，如果是，则返回当前值并将值加1，否则，抛出StopIteration异常，结束遍历。最后，我们创建了一个CustomIterator的实例，并使用for循环遍历它。

创建生成器

生成器实际上是一种特殊的迭代器，通过定义一个包含yield关键字的函数即可创建生成器。yield关键字用于在函数执行过程中返回一个值，并将控制权交回给调用者。当再次调用生成器时，它会从上次返回的位置继续执行，直到再次遇到yield。通过yield，生成器函数可以逐步产生值，而不需要一次性计算并返回所有值，节省了内存空间。与普通函数不同的是，生成器是一个返回迭代器的函数，只能用于迭代操作。调用一个生成器函数，返回的是一个迭代器对象。在下面的示例代码中，我们定义了get_odd生成器，用于生成小于num的奇数的迭代器。

def get_odd(num):
    ori = 1
    while ori < num:
        yield ori
        ori += 2

生成器在处理大量数据或需要按需生成数据的场景中，是非常有用的。比如：在处理文本文件时，我们可能不需要将整个文件一次性加载到内存中，而是可以使用生成器逐行读取文件。另外，在机器学习、大数据处理等领域，生成器也能够发挥重要作用。

Python中的生成器具有以下几个优点：

1、按需生成数据，有效节省内存空间。

2、能够处理大量数据，而不会导致内存溢出。

3、可以使用简单的代码实现复杂的迭代逻辑。

使用生成器

创建好生成器之后，我们就可以像使用迭代器一样使用生成器了。以上面的get_odd生成器为例，如果我们需要输出10以下的奇数，既可以使用next()函数，也可以使用for语句，示例代码如下。

def get_odd(num):
    ori = 1
    while ori < num:
        yield ori
        ori += 2
odd_generator = get_odd(10)
# 输出：1
print(next(odd_generator))
# 输出：3
print(next(odd_generator))
# 依次输出：5 7 9
for item in odd_generator:
    print(item)

考虑下面的应用场景：我们需要从文件中读取大量数据，并进行相应的处理。如果使用传统的列表或元组，可能会占用大量内存。此时，可以使用生成器逐行读取文件，从而有效节省内存。具体如何使用，可参考下面的示例代码。

def read_file(file_path):
    with open(file_path, 'r', encoding='utf-8', errors='ignore') as file:
        while line := file.readline():
            yield line.strip()
generator = read_file('./test.py')
# 使用生成器读取文件，并按行输出文件内容
for line in generator:
    print(line)