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