MK
摩柯社区 - 一个极简的技术知识社区
AI 面试

Python中的迭代器与生成器

2022-02-131.2k 阅读

Python 中的迭代器

在 Python 编程中,迭代是一种重要的概念,它允许我们遍历容器(如列表、元组、字典等)中的元素。迭代器(Iterator)是实现迭代功能的对象,它提供了一种在容器中逐个访问元素的方式,而无需一次性将所有元素加载到内存中。这在处理大型数据集时非常有用,因为它可以显著减少内存的使用。

可迭代对象(Iterable)

在理解迭代器之前,我们先来了解一下可迭代对象。可迭代对象是指那些可以被迭代的对象,比如列表、元组、字符串、字典等。在 Python 中,一个对象如果实现了 __iter__() 方法,那么它就是可迭代对象。这个方法返回一个迭代器对象,用于遍历该可迭代对象中的元素。

下面是一个简单的示例,展示如何判断一个对象是否为可迭代对象:

from collections.abc import Iterable

my_list = [1, 2, 3]
print(isinstance(my_list, Iterable))  # 输出: True

my_int = 5
print(isinstance(my_int, Iterable))  # 输出: False

在上述代码中,我们使用 isinstance() 函数和 collections.abc.Iterable 抽象基类来判断对象是否为可迭代对象。列表 my_list 是可迭代对象,而整数 my_int 不是。

迭代器协议

迭代器是遵循迭代器协议的对象。迭代器协议要求对象必须实现两个方法:__iter__()__next__()

  1. __iter__() 方法:这个方法返回迭代器对象本身,它使得迭代器对象可以被用于 for 循环等迭代环境中。
  2. __next__() 方法:这个方法返回迭代器的下一个元素。当没有更多元素可返回时,它会引发 StopIteration 异常,这表明迭代已经结束。

下面我们手动实现一个简单的迭代器:

class MyIterator:
    def __init__(self, limit):
        self.limit = limit
        self.current = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self.current >= self.limit:
            raise StopIteration
        self.current += 1
        return self.current - 1


my_iter = MyIterator(5)
for num in my_iter:
    print(num)

在上述代码中,我们定义了一个 MyIterator 类,它实现了迭代器协议。__init__() 方法初始化了迭代器的上限和当前值。__iter__() 方法返回迭代器本身,__next__() 方法返回下一个值,并在达到上限时引发 StopIteration 异常。我们可以使用 for 循环来遍历这个迭代器。

使用迭代器的优势

  1. 节省内存:迭代器允许我们逐个处理数据,而不是一次性加载所有数据到内存中。这在处理大型文件或数据流时非常有用,例如读取一个非常大的文本文件,一行一行地处理,而不是将整个文件读入内存。
with open('large_file.txt') as file:
    for line in file:
        # 处理每一行数据
        print(line.strip())

在上述代码中,open('large_file.txt') 返回的文件对象是一个迭代器,for 循环每次只读取文件的一行,大大节省了内存。 2. 延迟计算:迭代器可以在需要时才生成数据,而不是预先计算好所有数据。这对于一些计算量较大的数据生成非常有帮助。例如,生成斐波那契数列的迭代器:

class FibonacciIterator:
    def __init__(self):
        self.prev = 0
        self.curr = 1

    def __iter__(self):
        return self

    def __next__(self):
        result = self.prev
        self.prev, self.curr = self.curr, self.prev + self.curr
        return result


fib_iter = FibonacciIterator()
for i in range(10):
    print(next(fib_iter))

在这个斐波那契数列迭代器中,每次调用 __next__() 方法才计算并返回下一个斐波那契数,而不是预先计算好整个数列。

Python 中的生成器

生成器(Generator)是一种特殊的迭代器,它的创建更加简洁和高效。生成器使用 yield 关键字来暂停和恢复函数的执行,从而实现迭代器的功能。

生成器函数

生成器函数是一种定义生成器的方式。它看起来像一个普通函数,但使用 yield 关键字而不是 return。当生成器函数被调用时,它返回一个生成器对象,但不会立即执行函数体。只有当调用生成器对象的 __next__() 方法(或使用 next() 内置函数)时,函数体才会开始执行,直到遇到 yield 语句。yield 语句会暂停函数的执行,并返回一个值,下次调用 __next__() 方法时,函数会从暂停的地方继续执行。

下面是一个简单的生成器函数示例:

def simple_generator():
    yield 1
    yield 2
    yield 3


gen = simple_generator()
print(next(gen))  # 输出: 1
print(next(gen))  # 输出: 2
print(next(gen))  # 输出: 3
# print(next(gen))  # 这行代码会引发 StopIteration 异常

在上述代码中,simple_generator() 是一个生成器函数,它使用 yield 语句返回三个值。每次调用 next(gen) 时,函数执行到 yield 语句,返回对应的值,并暂停执行。当没有更多的 yield 语句时,再次调用 next() 会引发 StopIteration 异常。

生成器表达式

除了生成器函数,Python 还提供了生成器表达式来创建生成器。生成器表达式类似于列表推导式,但使用圆括号而不是方括号。它的语法如下:

(expression for item in iterable if condition)

生成器表达式返回一个生成器对象,而不是一个列表。这使得它在处理大型数据集时更加高效,因为它不会一次性生成所有数据。

下面是一个生成器表达式的示例:

gen_expression = (i * i for i in range(5))
for num in gen_expression:
    print(num)

在上述代码中,(i * i for i in range(5)) 是一个生成器表达式,它生成 0 到 4 的平方值。我们可以使用 for 循环来遍历这个生成器。

生成器的优势

  1. 简洁高效:生成器的语法更加简洁,尤其是使用生成器表达式时。与手动实现迭代器相比,生成器可以用更少的代码实现相同的功能。
  2. 节省内存:和迭代器一样,生成器也是按需生成数据,不会一次性占用大量内存。这在处理大型数据集或无限序列时非常有优势。例如,生成无限的奇数序列:
def odd_numbers():
    num = 1
    while True:
        yield num
        num += 2


odd_gen = odd_numbers()
for _ in range(10):
    print(next(odd_gen))

在上述代码中,odd_numbers() 是一个生成无限奇数序列的生成器函数。由于它是按需生成数据,所以不会占用过多内存。

迭代器与生成器的比较

  1. 实现方式
    • 迭代器:需要手动实现 __iter__()__next__() 方法来遵循迭代器协议。这在实现复杂的迭代逻辑时可能会比较繁琐。
    • 生成器:通过生成器函数(使用 yield 关键字)或生成器表达式来创建,实现更加简洁。生成器函数会自动创建一个遵循迭代器协议的对象。
  2. 内存使用
    • 迭代器:逐个返回数据,在处理大型数据集时能有效节省内存。但手动实现迭代器时,需要注意管理好状态和内存。
    • 生成器:同样按需生成数据,内存使用效率高。生成器的简洁性使得它在处理内存敏感的任务时更加方便。
  3. 代码复杂度
    • 迭代器:手动实现迭代器可能需要编写较多的代码来管理迭代状态和异常处理,代码复杂度相对较高。
    • 生成器:生成器函数和表达式的代码更加简洁明了,降低了代码复杂度,提高了可读性。

迭代器与生成器的应用场景

  1. 数据处理
    • 在处理大型文件或数据流时,无论是迭代器还是生成器都非常有用。例如,处理日志文件,一行一行地读取并分析日志数据:
def process_log(log_file):
    with open(log_file) as file:
        for line in file:
            # 分析日志行
            parts = line.split()
            if len(parts) >= 3:
                yield parts[0], parts[1], parts[2]


log_gen = process_log('app.log')
for date, time, message in log_gen:
    print(f'[{date} {time}] {message}')

在上述代码中,process_log() 函数是一个生成器函数,它逐行读取日志文件并解析出日期、时间和消息部分,然后通过 yield 返回。这样可以高效地处理大型日志文件,而不会占用过多内存。 2. 数据生成 - 生成器非常适合生成无限序列或大量数据的场景。比如生成随机数序列:

import random


def random_number_generator():
    while True:
        yield random.randint(1, 100)


rand_gen = random_number_generator()
for _ in range(5):
    print(next(rand_gen))

在这个示例中,random_number_generator() 是一个生成无限随机数序列的生成器函数。每次调用 next(rand_gen) 都会生成一个新的随机数。 3. 协同程序(Coroutine) - 生成器在协同程序中有重要应用。通过 yieldsend() 方法,可以实现多个任务之间的协作式多任务处理。例如:

def consumer():
    while True:
        value = yield
        print(f'Consumed: {value}')


def producer(consumer_obj):
    data = [1, 2, 3, 4, 5]
    for item in data:
        consumer_obj.send(item)


c = consumer()
next(c)  # 启动生成器
producer(c)

在上述代码中,consumer() 是一个生成器函数,它通过 yield 暂停并接收 producer() 发送的值。producer() 函数通过 send() 方法向 consumer() 发送数据,实现了简单的协同程序功能。

高级特性:生成器的 send()throw()close() 方法

  1. send() 方法
    • send() 方法不仅可以获取生成器的下一个值,还可以向生成器发送数据。它会将数据作为 yield 表达式的值传入生成器函数中。例如:
def generator_with_send():
    value = yield 1
    yield value


gen = generator_with_send()
print(next(gen))  # 输出: 1
print(gen.send(2))  # 输出: 2

在上述代码中,generator_with_send() 是一个生成器函数。第一次调用 next(gen) 时,函数执行到第一个 yield 1 语句,返回 1 并暂停。第二次调用 gen.send(2) 时,2 被赋值给 value,然后函数继续执行到第二个 yield value 语句,返回 2 并暂停。 2. throw() 方法 - throw() 方法用于在生成器内部引发异常。它允许我们在生成器外部向生成器内部抛出特定的异常,生成器可以捕获并处理这个异常。例如:

def generator_with_throw():
    try:
        yield 1
    except ValueError as ve:
        print(f'Caught ValueError: {ve}')
        yield 'Handled'


gen = generator_with_throw()
print(next(gen))  # 输出: 1
print(gen.throw(ValueError('Test error')))  # 输出: Handled

在上述代码中,generator_with_throw() 是一个生成器函数。第一次调用 next(gen) 时,函数执行到 yield 1 语句,返回 1 并暂停。第二次调用 gen.throw(ValueError('Test error')) 时,在生成器内部引发 ValueError 异常,生成器捕获并处理这个异常,然后返回 'Handled'。 3. close() 方法 - close() 方法用于关闭生成器。调用 close() 方法后,生成器不能再产生值,再次调用 next()send() 方法会引发 StopIteration 异常。例如:

def simple_generator_close():
    yield 1
    yield 2
    yield 3


gen = simple_generator_close()
print(next(gen))  # 输出: 1
gen.close()
# print(next(gen))  # 这行代码会引发 StopIteration 异常

在上述代码中,simple_generator_close() 是一个生成器函数。调用 gen.close() 后,生成器被关闭,再次调用 next(gen) 会引发 StopIteration 异常。

迭代器与生成器的嵌套使用

在实际编程中,我们经常会遇到需要嵌套使用迭代器和生成器的情况。例如,有一个包含多个列表的列表,我们想要逐个遍历每个子列表中的元素:

lists = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]


def nested_generator(lists):
    for sub_list in lists:
        for item in sub_list:
            yield item


nested_gen = nested_generator(lists)
for num in nested_gen:
    print(num)

在上述代码中,nested_generator() 是一个生成器函数,它嵌套了两个 for 循环,分别遍历外层列表和内层列表,通过 yield 逐个返回元素。我们可以使用 for 循环遍历这个嵌套生成器。

迭代器与生成器在 Python 标准库中的应用

  1. itertools 模块
    • itertools 模块提供了许多用于处理迭代器的函数,其中很多都使用了生成器。例如,count() 函数返回一个从指定值开始的无限迭代器:
import itertools


count_gen = itertools.count(1)
for _ in range(5):
    print(next(count_gen))

在上述代码中,itertools.count(1) 返回一个从 1 开始的无限迭代器,我们可以通过 next() 函数获取下一个值。 - chain() 函数可以将多个可迭代对象连接成一个迭代器:

list1 = [1, 2, 3]
list2 = [4, 5, 6]
chain_gen = itertools.chain(list1, list2)
for num in chain_gen:
    print(num)

在这个示例中,itertools.chain(list1, list2)list1list2 连接成一个迭代器,我们可以通过 for 循环遍历这个连接后的迭代器。 2. zip() 函数 - zip() 函数返回一个迭代器,它将多个可迭代对象中的元素一一配对。例如:

names = ['Alice', 'Bob', 'Charlie']
ages = [25, 30, 35]
zip_gen = zip(names, ages)
for name, age in zip_gen:
    print(f'{name} is {age} years old')

在上述代码中,zip(names, ages) 返回一个迭代器,将 namesages 中的元素一一配对。我们可以通过 for 循环遍历这个迭代器,获取配对后的结果。

总结迭代器与生成器的要点

  1. 迭代器:是遵循迭代器协议的对象,通过 __iter__()__next__() 方法实现迭代功能。它允许我们逐个访问容器中的元素,节省内存。手动实现迭代器需要较多代码来管理状态和异常。
  2. 生成器:是一种特殊的迭代器,通过生成器函数(使用 yield 关键字)或生成器表达式创建。生成器代码简洁,同样按需生成数据,节省内存。生成器还具有 send()throw()close() 等方法,可实现更复杂的功能。
  3. 应用场景:迭代器和生成器在数据处理、数据生成和协同程序等场景中都有广泛应用。在处理大型数据集或无限序列时,它们能显著提高内存使用效率。

在 Python 编程中,深入理解迭代器和生成器的概念和使用方法,可以帮助我们编写出更高效、简洁的代码,尤其是在处理大数据和复杂的迭代逻辑时。通过合理运用迭代器和生成器,我们可以充分发挥 Python 的优势,提升程序的性能和可维护性。