Python 自定义迭代器类

自定义迭代器类通过实现迭代器协议创建可迭代对象，精确控制迭代过程。

迭代器协议

迭代器必须实现两个方法：

• __iter__：返回迭代器对象自身
• __next__：返回下一个元素，无元素时抛出 StopIteration

基本示例：计数器

class Counter:
    "自定义计数迭代器"

    def __init__(self, start, end, step=1):
        self.current = start
        self.end = end
        self.step = step

    def __iter__(self):
        return self  # 返回自身作为迭代器

    def __next__(self):
        if self.current >= self.end:
            raise StopIteration  # 结束迭代
        value = self.current
        self.current += self.step
        return value


# 使用
counter = Counter(0, 10, 2)
for num in counter:
    print(num)  # 0, 2, 4, 6, 8

# 手动迭代
counter = Counter(5, 8)
print(next(counter))  # 5
print(next(counter))  # 6
print(next(counter))  # 7
print(next(counter))  # StopIteration

可重复迭代器

实现独立迭代器，支持多次迭代：

class Range:
    "模拟 range 的可重复迭代器"

    def __init__(self, start, end, step=1):
        self.start = start
        self.end = end
        self.step = step

    def __iter__(self):
        # 返回新迭代器对象，而非自身
        return RangeIterator(self.start, self.end, self.step)


class RangeIterator:
    "独立迭代器"

    def __init__(self, start, end, step):
        self.current = start
        self.end = end
        self.step = step

    def __iter__(self):
        return self

    def __next__(self):
        if self.current >= self.end:
            raise StopIteration
        value = self.current
        self.current += self.step
        return value


# 可多次迭代
r = Range(0, 5)
print(list(r))  # [0, 1, 2, 3, 4]
print(list(r))  # [0, 1, 2, 3, 4]  再次迭代成功

带状态的迭代器

class Fibonacci:
    "斐波那契数列迭代器"

    def __init__(self, max_count):
        self.max_count = max_count
        self.count = 0
        self.prev = 0
        self.curr = 1

    def __iter__(self):
        return self

    def __next__(self):
        if self.count >= self.max_count:
            raise StopIteration
        self.count += 1
        if self.count == 1:
            return 0
        if self.count == 2:
            return 1
        value = self.prev + self.curr
        self.prev, self.curr = self.curr, value
        return value


fib = Fibonacci(10)
print(list(fib))  # [0, 1, 1, 2, 3, 5, 8, 13, 21, 34]

反向迭代器

class Reverse:
    "反向迭代器"

    def __init__(self, data):
        self.data = data
        self.index = len(data)

    def __iter__(self):
        return self

    def __next__(self):
        if self.index <= 0:
            raise StopIteration
        self.index -= 1
        return self.data[self.index]


rev = Reverse([1, 2, 3, 4, 5])
print(list(rev))  # [5, 4, 3, 2, 1]

文件行迭代器

class FileLines:
    "文件行迭代器"

    def __init__(self, filename):
        self.file = open(filename, 'r')
        self.closed = False

    def __iter__(self):
        return self

    def __next__(self):
        line = self.file.readline()
        if line == '':
            self.file.close()
            self.closed = True
            raise StopIteration
        return line.strip()

    def __del__(self):
        if not self.closed:
            self.file.close()


# 使用
for line in FileLines('data.txt'):
    print(line)

无限迭代器

class InfiniteCounter:
    "无限计数迭代器"

    def __init__(self, start=0, step=1):
        self.current = start
        self.step = step

    def __iter__(self):
        return self

    def __next__(self):
        value = self.current
        self.current += self.step
        return value  # 永不抛出 StopIteration


# 必须手动限制
from itertools import islice
counter = InfiniteCounter(10, 5)
print(list(islice(counter, 5)))  # [10, 15, 20, 25, 30]

迭代器与可迭代对象

# 可迭代对象：只实现 __iter__
class IterableObject:
    def __init__(self, data):
        self.data = data

    def __iter__(self):
        return IteratorObject(self.data)


# 迭代器：实现 __iter__ 和 __next__
class IteratorObject:
    def __init__(self, data):
        self.data = data
        self.index = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self.index >= len(self.data):
            raise StopIteration
        value = self.data[self.index]
        self.index += 1
        return value


obj = IterableObject([1, 2, 3])
for item in obj:
    print(item)

与生成器对比

完整示例：树遍历迭代器

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


class TreeIterator:
    "二叉树遍历迭代器"

    def __init__(self, root, mode='inorder'):
        self.root = root
        self.mode = mode
        self.stack = []
        self._setup()

    def _setup(self):
        if self.mode == 'inorder':
            self._push_left(self.root)
        elif self.mode == 'preorder':
            self.stack.append(self.root)

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

    def __iter__(self):
        return self

    def __next__(self):
        if not self.stack:
            raise StopIteration

        node = self.stack.pop()

        if self.mode == 'inorder':
            self._push_left(node.right)
        elif self.mode == 'preorder' and node.right:
            self.stack.append(node.right)

        return node.value


# 使用
root = TreeNode(4,
    TreeNode(2, TreeNode(1), TreeNode(3)),
    TreeNode(6, TreeNode(5), TreeNode(7)))

print(list(TreeIterator(root, 'inorder')))   # [1, 2, 3, 4, 5, 6, 7]
print(list(TreeIterator(root, 'preorder')))  # [4, 2, 1, 3, 6, 5, 7]

要点总结

迭代器类适合复杂迭代逻辑、状态管理需求场景，简单场景优先使用生成器函数。

D:\git2\jwdev\articles\PYTHON\进阶\迭代器与生成器\自定义迭代器类.md