實現迭代器協議

問題

你想構建一個能支持迭代操作的自定義對象，并希望找到一個能實現迭代協議的簡單方法。

解法

目前為止，在一個對象上實現迭代最簡單的方式是使用一個生成器函數。在 4.2 小節中，使用 Node 類來表示樹形數據結構。你可能想實現一個以深度優先方式遍歷樹形節點的生成器。下面是代碼示例：

class Node:def __init__(self, value):self._value = valueself._children = []def __repr__(self):return 'Node({!r})'.format(self._value)def add_child(self, node):self._children.append(node)def __iter__(self):return iter(self._children)def depth_first(self):yield selffor c in self:yield from c.depth_first()# Example if __name__ == '__main__':root = Node(0)child1 = Node(1)child2 = Node(2)root.add_child(child1)root.add_child(child2)child1.add_child(Node(3))child1.add_child(Node(4))child2.add_child(Node(5))for ch in root.depth_first():print(ch)# Outputs Node(0), Node(1), Node(3), Node(4), Node(2), Node(5)

在這段代碼中， depth_first() 方法簡單直觀。它首先返回自己本身并迭代每一個子節點并通過調用子節點的 depth_first() 方法 (使用 yield from 語句) 返回對應元素。

討論

Python 的迭代協議要求一個__iter__() 方法返回一個特殊的迭代器對象，這個迭代器對象實現了__next__() 方法并通過 StopIteration 異常標識迭代的完成。但是，實現這些通常會比較繁瑣。下面我們演示下這種方式，如何使用一個關聯迭代器類重新實現 depth_first() 方法：

class Node2:def __init__(self, value):self._value = valueself._children = []def __repr__(self):return 'Node({!r})'.format(self._value)def add_child(self, node):self._children.append(node)def __iter__(self):return iter(self._children)def depth_first(self):return DepthFirstIterator(self)class DepthFirstIterator(object):'''Depth-first traversal'''def __init__(self, start_node):self._node = start_nodeself._children_iter = Noneself._child_iter = Nonedef __iter__(self):return selfdef __next__(self):# Return myself if just started; create an iterator for childrenif self._children_iter is None:self._children_iter = iter(self._node)return self._node# If processing a child, return its next itemelif self._child_iter:try:nextchild = next(self._child_iter)return nextchildexcept StopIteration:self._child_iter = Nonereturn next(self)# Advance to the next child and start its iterationelse:self._child_iter = next(self._children_iter).depth_first()return next(self)

DepthFirstIterator 類和上面使用生成器的版本工作原理類似，但是它寫起來很繁瑣，因為迭代器必須在迭代處理過程中維護大量的狀態信息。坦白來講，沒人愿意寫這么晦澀的代碼。將你的迭代器定義為一個生成器后一切迎刃而解。

總結

以上是生活随笔為你收集整理的《Python Cookbook 3rd》笔记（4.4）：实现迭代器协议的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。