實(shí)現(xiàn)迭代器協(xié)議

問題

你想構(gòu)建一個(gè)能支持迭代操作的自定義對象，并希望找到一個(gè)能實(shí)現(xiàn)迭代協(xié)議的簡單方法。

解法

目前為止，在一個(gè)對象上實(shí)現(xiàn)迭代最簡單的方式是使用一個(gè)生成器函數(shù)。在 4.2 小節(jié)中，使用 Node 類來表示樹形數(shù)據(jù)結(jié)構(gòu)。你可能想實(shí)現(xiàn)一個(gè)以深度優(yōu)先方式遍歷樹形節(jié)點(diǎn)的生成器。下面是代碼示例：

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() 方法簡單直觀。它首先返回自己本身并迭代每一個(gè)子節(jié)點(diǎn)并通過調(diào)用子節(jié)點(diǎn)的 depth_first() 方法 (使用 yield from 語句) 返回對應(yīng)元素。

討論

Python 的迭代協(xié)議要求一個(gè)__iter__() 方法返回一個(gè)特殊的迭代器對象，這個(gè)迭代器對象實(shí)現(xiàn)了__next__() 方法并通過 StopIteration 異常標(biāo)識迭代的完成。但是，實(shí)現(xiàn)這些通常會比較繁瑣。下面我們演示下這種方式，如何使用一個(gè)關(guān)聯(lián)迭代器類重新實(shí)現(xiàn) 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 類和上面使用生成器的版本工作原理類似，但是它寫起來很繁瑣，因?yàn)榈鞅仨氃诘幚磉^程中維護(hù)大量的狀態(tài)信息。坦白來講，沒人愿意寫這么晦澀的代碼。將你的迭代器定義為一個(gè)生成器后一切迎刃而解。

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。