Search code examples
pythoniteratorgeneratoryield

Why does the iteration stop at the first yield?


I'm trying to walk an AST and yield all the matching nodes. My iteration stops at the first match, instead of going trough all the nodes. What am I missing?

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
                # why doesn't the execution continue from here?
            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)

As you can see, the result returns the first match only.

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

Solution

  • As Karl Knechtel answered before, the error lies in the test function:

    result = []
    while result.append(next(search_comment(test_ast))):
        pass
    

    changed to

    result = list(search_comment(test_ast))