4.4 實現迭代器協議 · Python3 CookBook中文版

## 問題你想構建一個能支持迭代操作的自定義對象，并希望找到一個能實現迭代協議的簡單方法。 ## 解決方案目前為止，在一個對象上實現迭代最簡單的方式是使用一個生成器函數。在4.2小節中，使用Node類來表示樹形數據結構。你可能想實現一個以深度優先方式遍歷樹形節點的生成器。下面是代碼示例： class Node: def __init__(self, value): self._value = value self._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 self for 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 = value self._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_node self._children_iter = None self._child_iter = None def __iter__(self): return self def __next__(self): # Return myself if just started; create an iterator for children if self._children_iter is None: self._children_iter = iter(self._node) return self._node # If processing a child, return its next item elif self._child_iter: try: nextchild = next(self._child_iter) return nextchild except StopIteration: self._child_iter = None return next(self) # Advance to the next child and start its iteration else: self._child_iter = next(self._children_iter).depth_first() return next(self) `DepthFirstIterator` 類和上面使用生成器的版本工作原理類似，但是它寫起來很繁瑣，因為迭代器必須在迭代處理過程中維護大量的狀態信息。坦白來講，沒人愿意寫這么晦澀的代碼。將你的迭代器定義為一個生成器后一切迎刃而解。