遍历AST时迭代在第一个yield处停止，如何获取所有匹配节点？

问题描述 #

尝试遍历AST并生成所有匹配Comment类型的节点路径，但迭代仅在第一个匹配项处停止，无法获取全部结果。以下是代码及测试结果：

测试代码 #

import unittest

test_ast = {
    'body': [
        {
            'type': 'BlockComment',
            'value': 'foo'
        },
        {
            'type': 'LineComment',
            'value': ' bar'
        }
    ]
}


def search_comment(haystack, path=()):
    if isinstance(haystack, list):
        for idx, item in enumerate(haystack):
            yield from search_comment(item, path=path + (idx,))
    elif isinstance(haystack, dict):
        for key, value in haystack.items():
            if key == 'type' and 'Comment' in value:
                yield path
                # 为什么执行到这里就不再继续了？
            yield from search_comment(value, path=path + (key,))


class IteratorTests(unittest.TestCase):
    def test_search_comment(self):
        expected = [('body', 0), ('body', 1)]
        result = []
        while result.append(next(search_comment(test_ast))):
            pass
        self.assertEqual(expected, result)

测试结果 #

Expected :[('body', 0), ('body', 1)]
Actual   :[('body', 0)]

问题原因 #

1. 测试代码的致命错误 #

在测试方法中，每次调用search_comment(test_ast)都会创建一个全新的生成器实例，而非复用同一个生成器。因此每次next()调用都只会获取第一个匹配项，当result.append()返回None时，while循环直接终止，导致只收集到第一个结果。

2. 遍历逻辑的冗余（非直接原因，但需优化） #

当字典的key为type且值包含Comment时，yield path之后代码还会递归处理value（即字符串类型的BlockComment/LineComment），这部分递归不会产生任何有效结果，属于冗余操作。

修复方案 #

修复测试代码 #

先创建一个生成器实例，再循环遍历该实例的所有结果：

def test_search_comment(self):
    expected = [('body', 0), ('body', 1)]
    result = []
    comment_generator = search_comment(test_ast)
    for path in comment_generator:
        result.append(path)
    self.assertEqual(expected, result)

或更简洁的写法：

def test_search_comment(self):
    expected = [('body', 0), ('body', 1)]
    result = list(search_comment(test_ast))
    self.assertEqual(expected, result)

优化遍历函数（可选） #

去掉冗余的递归调用，匹配到Comment类型后无需再递归处理type对应的字符串值：

def search_comment(haystack, path=()):
    if isinstance(haystack, list):
        for idx, item in enumerate(haystack):
            yield from search_comment(item, path=path + (idx,))
    elif isinstance(haystack, dict):
        for key, value in haystack.items():
            if key == 'type' and 'Comment' in value:
                yield path
                # 跳过对字符串value的递归，无有效子节点
                continue
            yield from search_comment(value, path=path + (key,))

验证结果 #

修复后运行测试，result会正确收集到两个匹配路径：[('body', 0), ('body', 1)]，与预期一致。

内容的提问来源于stack exchange，提问作者Christian P