使用Python生成器进行先序遍历，并带有忽略子树的机制

Question

我编写了以下代码来对Python dict进行预排序遍历，其中可能包含其他dict：

def preorder_traversal(obj):
    yield obj
    if isinstance(obj, dict):
        for k, v in obj.iteritems():
            for o in preorder_traversal(v):
                yield o

下面是一些它的行为示例：

^{pr2}$

树可能会变得非常大，所以我想添加一种机制，通过这种机制，pre-order遍历的使用者可以中止对整个子树的搜索。在

这是我想出的代码，包括一个鼻子测试用例。测试失败，如下所述。在

class IgnoreSubtree(Exception):
    pass    

def preorder_traversal(obj):
    try:
        yield obj
    except IgnoreSubtree:
        return

    if isinstance(obj, dict):
        for k, v in obj.iteritems():
            iterator = preorder_traversal(v)
            for o in iterator:
                try:
                    yield o
                except IgnoreSubtree as e:
                    try:
                        iterator.throw(e)
                    except StopIteration:  # WHY?
                        pass

def test_ignore_subtree():
    obj = {'a': {'c': 3}, 'b': 2}
    iterator = preorder_traversal(obj)
    out = []
    for o in iterator:
        out.append(o)
        if o == {'c': 3}:
            iterator.throw(IgnoreSubtree)

    eq_([obj, {'c': 3}, 2], out)

测试失败，错误如下：

AssertionError: [{'a': {'c': 3}, 'b': 2}, {'c': 3}, 2] !=
                [{'a': {'c': 3}, 'b': 2}, {'c': 3}]

也就是说，IgnoreSubtree也中止了顶层对象中k/v对的迭代，它并没有仅仅删去{'c': 3}子树。在

这个密码怎么了？为什么在上面的注释位置抛出StopIteration？这是为这个函数实现子树剪枝的合理方法，还是有更好的方法来实现它？在

Answer 1

正如audiodude所提到的，您的iterator.throw(IgnoreSubtree)返回iterator的下一个值（暂时忽略了复杂的异常处理），因此它正在消耗您期望在test_ignore_subtree中的下一次循环的2附加到out。在

您还询问了抛出StopIteration的原因；Exception被抛出/捕获的顺序是：

• iterator.throw(IgnoreSubtree)抛出一个在preorder_traversal的内循环中捕获的IgnoreSubtree
• 使用IgnoreSubtree将{}路由到内部迭代器中
• IgnoreSubtree在except IgnoreSubtree:处被捕获，并调用return；但是，iterator.throw(e)希望从内部迭代器中获得下一个值，该迭代器刚刚进行了return，因此引发了一个StopIteration。在
• 在原始的iterator.throw(IgnoreSubtree)返回之前，它将再次通过preorder_traversal中的外部循环，因为它希望从iterator返回下一个值。在

我希望这有帮助！在

更新

以下是我将使用的基本方案的实现，以及通过nosetest：

from nose.tools import eq_

def preorder_traversal(obj, ignore_only_descendents_of=None, ignore_subtrees=None):
    if ignore_subtrees and obj in ignore_subtrees:
        return

    yield obj

    if ignore_only_descendents_of and obj in ignore_only_descendents_of:
        return

    if isinstance(obj, dict):
        for k, v in iter(sorted(obj.iteritems())):
            iterator = preorder_traversal(v, ignore_only_descendents_of, ignore_subtrees)
            for o in iterator:
                yield o


def test_ignore_subtree():
    obj = {'a': {'c': 3}, 'b': 2, 'd': {'e': {'f': 4}}, 'g': 5, 'h': 6}
    ignore_only_descendents_of = [{'e': {'f': 4}}]
    ignore_subtrees = [{'c': 3}, 5]

    iterator = preorder_traversal(obj, ignore_only_descendents_of, ignore_subtrees)
    out = []

    for o in iterator:
        out.append(o)

    expected = [obj, 2, {'e': {'f': 4}}, 6]
    eq_(expected, out)

注意事项：

• 您的示例允许排除{'c':3}的子代，同时包含{'c':3}本身；我发现这有点令人困惑，因为我希望您通常希望排除包括其根的整个子树，所以我将preorder_traversal更改为以每种方式排除两个可选的内容列表。在
• 将子树移到迭代器本身看起来更简洁；您可以避免使用Exceptions来完全控制流。在
• 更复杂的例子演示了两种类型的子树排除。在

Answer 2

首先，为什么要提出StopIteration。您对preorder_traversal的定义从以下开始：

try:
    yield obj
except IgnoreSubtree:
    return

在生成器中，return语句等价于raise StopIteration。在python3.3+中，您实际上可以使用return value，它相当于raise StopIteration(value)。在

因此，您正在某个异常中throw，它被执行return的生成器捕获，从而引发一个StopIteration。无论何时调用send，next或throw如果生成器在没有找到yield的情况下结束其执行，则可能会引发next或{}，因此每当跳过子树将结束迭代时，您在测试中使用的代码注定会引发一个StopIteration。在

换句话说，您的测试是有缺陷的，因为throw调用可以引发异常，即使您有正确的生成器实现。因此，您应该在try语句中包装该调用：

^{pr2}$

或者，您可以使用^{}上下文管理器来取消StopIteration：

with suppress(StopIteration):
    for o in iterator:
        ...
        iterator.throw(IgnoreSubtree)

如果您不使用python3.4，那么可以使用^{}修饰符轻松地重新实现这个上下文管理器（从python2.6开始就可以使用它）：

def suppress(*exceptions):
    try:
        yield
    except exceptions:
        pass

你的代码基本上是正确的。如果您使用的是python3.3+，可以将其简化为：

def preorder_traversal(obj):
    try:
        yield obj
    except IgnoreSubtree:
        return
    else:
        if isinstance(obj, dict):
            for k, v in obj.items():
                yield from preorder_traversal(v)

一旦外部的StopIteration被抑制，throw，您的实现不会对我产生任何错误。结果也是你所期望的。 不幸的是，如果没有yield from，我看不到任何简化控制流的方法。在

Answer 3

通过对迭代器的异常进行throw来实现子树修剪，会在生成器和使用它的函数中产生混乱、易出错的代码。看看这里的一些答案，我认为过滤器回调函数是一种更合理的方法。在

这是Ryan's answer的概括：

def preorder_traversal(obj, bailout_fn=None):
    yield obj
    if bailout_fn and bailout_fn(obj):
        return

    if isinstance(obj, dict):
        for k, v in obj.iteritems():
            for o in preorder_traversal(v, bailout_fn):
                yield o

下面是一些鼻子测试来演示它的使用方法：

^{pr2}$