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

問題

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

解決方案

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

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

討論

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