Search code examples
python-3.xrecursiongeneratoryieldyield-from

How to yield objects recursively using python generator?


I am writing a generator function which recursively walks through all the child nodes for a given astroid node.

In the example below, node is an astroid functiondef node. node.getchildren() returns a generator with subnodes in the node.

My goal is to yield every node contained. ( even in the subnode )

def recursive_walk(node):
    try:
        for subnode in list(node.get_children()):
            # yield subnode
            print(subnode)
            recursive_walk(subnode)            

    except AttributeError:
        # yield node       
        print(node)

    except TypeError:
        # yield node  
        print(node)

Here, if I have commented out the yield statement. For print statement, I am getting my desired result, but if I yield the node, I am not getting the desired output.

For reproducing this: - install astroid

import astroid

node = astroid.extract_node('''
def test_function(something): #@
    """Test function for getting subnodes"""
    assign_line = "String"
    ctx = len(assign_line)
    if ctx > 0:
        if ctx == 1:
            return 1
        if ctx > 10:
            return "Ten"
    else:
        return 0
''')

Solution

  • def recursive_walk(node):
        """
        Generator function to recursively walk through a given node.
    
        Yields every astroid node inside.
        """
        try:
            for subnode in node.get_children():
                yield subnode
                yield from recursive_walk(subnode)
    
        except (AttributeError, TypeError):
            yield node
    

    This does the work.