How to yield objects recursively using python generator?

Question

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
''')

I have this error `Traceback (most recent call last): File "main.py", line 5, in node = astroid.extract_node(''' AttributeError: module 'astroid' has no attribute 'extract_node'` — Netwave, Jun 03 '19 at 12:04
Strange. See the last section [here.](http://pylint.pycqa.org/projects/astroid/en/latest/inference.html) — Srijan Saurav, Jun 04 '19 at 10:00

score 2 · Accepted Answer · edited Feb 21 '21 at 08:42

2

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.

edited Feb 21 '21 at 08:42

aviso

2,371
1
14
15

answered Jun 04 '19 at 10:03

Srijan Saurav

81
5

score 1 · Answer 2 · answered Jun 03 '19 at 10:54

1

You can recursively map the function to the childs and use yield from:

def recursive_walk(node):
    # yield the base node
    yield node
    # for each of the recursive calls yield the remaining ones
    yield from map(recursive_walk, node.get_children())

In case they have no children it will just yield nothing and continue with the next ones.

answered Jun 03 '19 at 10:54

Netwave

40,134
6
50
93

@SrijanSaurav, since you are new, you may consider reading this: https://stackoverflow.com/help/someone-answers , Happy coding! – Netwave Jun 03 '19 at 11:18
Sure. One follow up question: In the output, am getting the base node, and other generator objects, which is not what I need. Is there a way to retrieve all the nodes from those generator objects as well, so that I can have just one function which generates all the nodes ? – Srijan Saurav Jun 03 '19 at 11:32
you just need to call to `list(recursice_walk(base_node))` to retrieve all of them. – Netwave Jun 03 '19 at 11:33
Output: `[, , , , , ]` – Srijan Saurav Jun 03 '19 at 11:40
Also, the generator objects contain other generators for its nested nodes. – Srijan Saurav Jun 03 '19 at 11:42
Uhm that is weird, are you using `yield_from`?, you can try `list(itertools.chain.from_iterable(recursive_walk(base_node)))` – Netwave Jun 03 '19 at 11:42
using `yield from` The same changes you suggested. And with itertools, getting this : `['self', 'tpl', 'ctx', , , , , , , , , , , ]` – Srijan Saurav Jun 03 '19 at 11:49
1

Let me Edit my question and add some more details for you to reproduce. – Srijan Saurav Jun 03 '19 at 11:50

How to yield objects recursively using python generator?

2 Answers2