如何在自定义 Python 对象中实现迭代

简介

Python 的内置迭代机制，如 for 循环和列表推导式，是处理数据集合的强大工具。然而，如果你想创建自己的可迭代自定义对象会怎样呢？在本教程中，我们将探讨如何在自定义 Python 对象中实现迭代，让你能够充分利用 Python 迭代功能的全部威力。

理解 Python 中的迭代

什么是迭代？

编程中的迭代是指重复执行一组指令或一段代码的过程。在 Python 中，迭代是一个基本概念，它允许你处理序列（如列表、元组和字符串）以及其他可迭代对象。

可迭代对象

可迭代对象是指可以被迭代的对象，这意味着它可以被循环遍历，并且其元素可以逐个访问。在 Python 中，常见的可迭代对象包括：

列表
元组
字符串
字典
集合
文件
实现迭代器协议的自定义对象

`for` 循环

for 循环是在 Python 中迭代可迭代对象最常用的方式。for 循环允许你为可迭代对象中的每个元素执行一段代码。

fruits = ['apple', 'banana', 'cherry']
for fruit in fruits:
    print(fruit)

输出：

apple
banana
cherry

`while` 循环

while 循环是在 Python 中实现迭代的另一种方式。只要指定的条件为真，while 循环就会继续执行一段代码。

count = 0
while count < 5:
    print(count)
    count += 1

输出：

迭代器与迭代器协议

在幕后，Python 中的 for 循环和其他迭代机制使用迭代器协议。迭代器是实现迭代器协议的对象，该协议定义了两个方法：__iter__() 和 __next__()。__iter__() 方法返回迭代器对象本身，__next__() 方法返回序列中的下一个元素。

class MyIterator:
    def __init__(self, data):
        self.data = data
        self.index = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self.index < len(self.data):
            result = self.data[self.index]
            self.index += 1
            return result
        else:
            raise StopIteration()

my_iterator = MyIterator([1, 2, 3, 4, 5])
for item in my_iterator:
    print(item)

输出：

实现自定义迭代器

迭代器协议

如前所述，Python 中的迭代器协议定义了两个方法：__iter__() 和 __next__()。要创建自定义迭代器，你需要在自己的类中实现这两个方法。

实现 `iter()` 方法

__iter__() 方法应该返回迭代器对象本身。当你使用 iter() 函数或在 for 循环中使用该对象时，会调用此方法。

class MyIterator:
    def __init__(self, data):
        self.data = data
        self.index = 0

    def __iter__(self):
        return self

实现 `next()` 方法

__next__() 方法应该返回序列中的下一个元素。如果没有更多元素，它应该引发 StopIteration 异常。

class MyIterator:
    def __init__(self, data):
        self.data = data
        self.index = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self.index < len(self.data):
            result = self.data[self.index]
            self.index += 1
            return result
        else:
            raise StopIteration()

使用自定义迭代器

一旦你实现了迭代器协议，就可以在 for 循环或其他迭代机制中使用自定义迭代器。

my_iterator = MyIterator([1, 2, 3, 4, 5])
for item in my_iterator:
    print(item)

输出：

使用生成器进行延迟求值

生成器是一种特殊类型的函数，可用于创建自定义迭代器。生成器使用 yield 关键字一次返回一个值，而不是在内存中构建完整的列表。

def my_generator(n):
    i = 0
    while i < n:
        yield i
        i += 1

my_gen = my_generator(5)
for item in my_gen:
    print(item)

输出：

与创建完整列表相比，生成器在处理大型或无限数据集时可能更节省内存。

应用迭代自定义对象

自定义迭代器的用例

自定义迭代器在各种场景中都可能很有用，例如：

迭代大型或无限数据集而不消耗过多内存
实现可迭代的自定义数据结构
提供更直观或特定领域的方式来迭代数据

示例：迭代二叉搜索树

让我们考虑一个可以使用自定义迭代器进行迭代的二叉搜索树（BST）示例。

class Node:
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None

class BinarySearchTree:
    def __init__(self):
        self.root = None

    def insert(self, value):
        ## 为简洁起见，省略插入方法的实现

    def __iter__(self):
        return BSTIterator(self.root)

class BSTIterator:
    def __init__(self, root):
        self.stack = []
        self.push_left_children(root)

    def __next__(self):
        if not self.stack:
            raise StopIteration()
        node = self.stack.pop()
        self.push_left_children(node.right)
        return node.value

    def push_left_children(self, node):
        while node:
            self.stack.append(node)
            node = node.left

## 示例用法
bst = BinarySearchTree()
bst.insert(5)
bst.insert(3)
bst.insert(7)
bst.insert(1)
bst.insert(4)
bst.insert(6)
bst.insert(8)

for value in bst:
    print(value)

输出：

在这个示例中，我们为 BinarySearchTree 类实现了一个自定义迭代器。BSTIterator 类使用一个栈来对二叉搜索树进行中序遍历，使我们能够按排序顺序迭代树的元素。

迭代无限序列

自定义迭代器还可用于处理无限序列，例如斐波那契数列或质数序列。通过使用生成器，我们可以创建能够即时生成下一个元素的迭代器，而无需将整个序列存储在内存中。

def fibonacci_generator():
    a, b = 0, 1
    while True:
        yield a
        a, b = b, a + b

fib = fibonacci_generator()
for i in range(10):
    print(next(fib))

输出：

通过使用生成器函数，我们可以创建一个能够无限生成斐波那契数列的迭代器，而不会消耗大量内存。

总结

在本教程结束时，你将对如何在自己的自定义 Python 对象中实现迭代有扎实的理解。你将学习自定义迭代器的基础知识，如何将它们应用于你的对象，并探索这种强大技术的实际用例。掌握 Python 中的自定义迭代将使你能够创建更灵活、高效和直观的代码，这些代码能无缝集成到该语言的核心特性中。