Python 中的迭代器实现原理

在数学中，集合表示由一个或多个确定的元素所构成的整体。在 Python 中，列表、元组、集合可以用于表示数学中的集合。

例如，分别使用列表、元组、集合表示了一个包含 3 个字符串的集合：

• 列表 [‘www’, ‘imooc’, ‘com’]
• 元组 (‘www’, ‘imooc’, ‘com’)
• 集合 {‘www’, ‘imooc’, ‘com’}

1. 可迭代对象 iterable

1.1 什么是可迭代对象

Python 提供了 for … in 循环，用于对列表、元组、集合中的元素进行遍历。能够被 for … in 循环遍历的对象被称为可迭代对象 iterable，列表、元组、集合均属于可迭代对象。使用 for … in 循环遍历可迭代对象的例子如下：

• 遍历列表的代码

list = ['www', 'imooc', 'com']
for item in list:
    print(item)

• 遍历元组的代码

tuple = ('www', 'imooc', 'com')
for item in tuple:
    print(item)

• 遍历集合的代码

set = {'www', 'imooc', 'com'}
for item in set:
    print(item)

1.2 尽可能使用 for … in 循环进行遍历

如果需要遍历的对象是列表，可以通过访问索引的方式进行遍历，代码如下：

strings = ['www', 'imooc', 'com']
i = 0
while i < len(strings):
    string = strings[i]
    print(string)
    i = i + 1

• 在第 1 行，使用列表表示 strings
• 在第 3 行，通过 len(strings) 获取列表 strings 中字符串的数量
• 在第 4 行，通过 strings[i] 访问第 i 个元素

以上的遍历方式中，要求 strings 是一个列表，如果 strings 的数据结构发生变化：使用集合而不是列表表示 strings，那么通过访问索引的方式进行遍历的代码就会失效。

strings = {'www', 'imooc', 'com'}
i = 0
while i < len(strings):
    string = strings[i]
    print(string)
    i = i + 1

• 在第 1 行，使用集合表示 strings
• 在第 3 行，通过 len(strings) 获取集合 strings 中字符串的数量
• 在第 4 行，通过 strings[i] 访问第 i 个元素

因为 strings 是一个集合，不支持索引操作，会导致运行错误：

Traceback (most recent call last):
  File "strings.py", line 5, in <module>
    string = strings[i]
TypeError: 'set' object does not support indexing

应尽可能使用 for … in 循环遍历可迭代对象，如果可迭代对象的数据类型发生变化，从列表变成集合，使用for … in 循环遍历的代码则无需改变。

2. 迭代器 iterator

1.1 什么是迭代器

迭代器 iterator 是一个特殊的对象，用于遍历访问可迭代对象 iterable。Python 通过迭代器 iterator 实现 for … in 循环语句，用户编写的 for … in 循环代码如下：

for item in iterable:
    print(item)

这段 for … in 循环代码会被翻译为如下：

iterator = iter(iterable)
while True:
    try:
        item = next(iterator)
        print(item)
    except StopIteration:
        break

• 在第 1 行，内置函数 iter 获取可迭代对象 iterable 的迭代器 iterator
• 在第 4 行，内置函数 next 获取迭代器 iterator 返回的下一个元素
• 在第 6 行，当迭代器遍历完全部元素后，抛出一个特殊的异常 StopIteration，表示迭代结束

1.2 列表的迭代器

下面通过一个具体的例子，了解如何通过迭代器实现 for … in 循环，使用 for … in 循环遍历列表的代码如下：

list = ['www', 'imooc', 'com']
for item in list:
    print(item)

Python 把以上 for … in 循环转换为如下功能等价的代码：

list = ['www', 'imooc', 'com']
listIterator = iter(list)
while True:
    try:
        item = next(listIterator)
        print(item)
    except StopIteration:
        break

以上两段代码均输出相同的结果，如下所示：

www
imooc
com

3. 迭代协议

使用迭代器遍历访问可迭代对象，要求迭代器和可迭代对象遵循迭代协议，迭代协议如下：

1. 可迭代对象 iterable 提供成员方法 __iter__，该方法返回用于遍历的迭代器 iterator

class Iterable:
    def __iter__(self):

2. 迭代器 iterator 提供成员方法 __next__，该方法返回下一个被遍历的元素

class Iterator:
    def __next__(self):

3. 异常 StopIteration，当遍历完全部的元素后，成员方法 __next__ 抛出一个特殊的异常 Stop Iteration 表示遍历结束

4. 内置函数 iter，用于获取可迭代对象对应的迭代器

def iter(iterable):
    iterator = iterable.__iter__()
    return iterator

• 在第 1 行，iter 的输入参数是可迭代对象 iterable
• 在第 2 行，调用成员方法 __iter__
• 在第 3 行，返回迭代器 iterator

5. 内置函数 next，用于获取下一个被遍历的元素

def next(iterator):
    item = iterator.__next__()
    return item

• 在第 1 行，next 的输入参数是迭代器 iterator
• 在第 2 行，调用成员方法 __next__
• 在第 3 行，返回被遍历的元素

根据以上的迭代协议，即可将 for … in 循环翻译为如下等价代码：

iterator = iter(iterable)
while True:
    try:
        item = next(iterator)
        print(item)
    except StopIteration:
        break

4. 实现一个自定义的迭代器

4.1 通过单链表实现堆栈

通过单链表实现堆栈，图示如下：

在上图中，每个节点有两个字段: item 和 next，item 用于存储数据，next 指向下一个节点，head 指针指向堆栈的顶部。描述堆栈的 Python 代码如下：

class Node:
    def __init__(self, item):
        self.item = item
        self.next = None

class Stack:
    def __init__(self):
        self.head = None

    def push(self, item):
        node = Node(item)
        node.next = self.head
        self.head = node

stack = Stack()
stack.push('a')
stack.push('b')
stack.push('c')

• 在第 1 行，定义了类 Node 用于描述链表中的节点
• 在第 6 行，定义了类 Stack 描述堆栈
  • 在第 8 行，定义了头指针 head，指向链表中的首个节点
  • 在第 10 行，定义了成员方法 push，将元素压如到堆栈中
    • 在第 11 行，创建一个新节点 node
    • 在第 12 行，新节点 node 的 next 指向头结点
    • 在第 13 行，头结点指向新节点
• 在第 15 行，创建一个对象 stack
• 在第 16 行到第 18 行，依次压入 3 个元素 ‘a’、‘b’、‘c’

4.2 实现迭代协议

class StackIterator:
    def __init__(self, stack):
        self.stack = stack
        self.cursor = self.stack.head

    def __next__(self):
        if self.cursor == None:
            raise StopIteration
        else:
            item = self.cursor.item
            self.cursor = self.cursor.next
            return item

• 在第 1 行，定义类 StackIterator
  • 类 Stack 是可迭代对象
  • 类 StackIterator 是迭代器
• 在第 2 行，定义构造函数，参数 stack 是被遍历的对象
  • 在第 4 行，成员变量 cursor 指向了当前正在遍历的元素，初始化被设置为链表的头结点
• 在第 6 行，定义方法 __next__
  • 在第 7 行，如果变量 cursor 等于 None，表示已经到达链表的尾部，即遍历完全部的元素了
    • 在第 8 行，抛出异常 StopIteration 表示遍历结束
  • 在第 9 行，如果变量 cursor 不等于 None
    • 在第 10 行，记录下当前正在遍历的元素
    • 在第 11 行，将 cursor 指向下一个元素

在定义了 StackIterator 后，在 Stack 中增加一个新的成员方法 __iter__，返回 Stack 对应的迭代器，代码如下：

class Stack:
    def __iter__(self):
        return StackIterator(self)   

4.3 通过 while 循环遍历堆栈

在实现了迭代协议后，使用 while 循环显示的使用 iter、next、StopIteration 完成对 stack 的遍历，代码如下：

stackIterator = iter(stack)
while True:
    try:
        item = next(stackIterator)
        print(item)
    except StopIteration:
        break

程序依次压入 ‘a’、‘b’、‘c’，遍历时以压入相反的顺序输出，结果如下：

c
b
a

4.4 通过 for … in 循环遍历堆栈

在实现了迭代协议后，可以通过 for … in 循环进行遍历，代码如下：

for item in stack:
    print(item)

与上一节的代码相比，代码要简洁很多，程序输出相同的结果如下：

c
b
a