How to find ast nodes and its sub nodes in Python

Question

I am trying to traverse through ast nodes in Python to fetch value node's Name ids for all Assign Nodes. Ast.walk() method gives me nodes randomly.

Assumption : Only for Assign nodes

for node in ast.walk(tree):


    if isinstance(node,(ast.Assign)):
        if isinstance(node.value,(ast.Name)):
            print('Assign Value:-','lineno: ',node.value.lineno,' id :',node.value.id)
        if isinstance(node.value,(ast.Tuple)):
            for i in range(len(node.value.elts)):
                if isinstance(node.value.elts[i],(ast.Name)):
                    print('Assign Value:-','lineno: ',node.value.elts[i].lineno,' id :',node.value.elts[i].id)

Here, sometimes ast gives BinOps, Call etc inside value node. I have to get Name.id inside value node.

Ex: a,b = (c+d), e --> Expected Values should be c,d,e..but I get only e . Also, how to tag (c,d) to target 'a' .

Target can be fetched from Assign Node .

Answer 1

You can walk the original ast, and for every ast.Assign, recursively extract target names and associated value ids:

import ast
def names(a):
   return [i.id for i in ast.walk(a) if isinstance(i, ast.Name)]

def bindings(target, value):
   if isinstance(target, ast.Name):
      yield [target.id, names(value)]
   elif isinstance(target, (ast.List, ast.Tuple)):
      if not isinstance(value, (ast.List, ast.Tuple)):
         yield [tuple(names(target)), names(value)]
      else:
         for a, b in zip(target.elts, value.elts):
            yield from bindings(a, b)

def assign_bindings(s):
   return [dict(bindings(i.targets[0], i.value)) for i in ast.walk(ast.parse(s)) 
            if isinstance(i, ast.Assign)]

The code above allows you to pass assign_bindings a source string that has assignment statements with unpacking, single name bindings, tuples, expressions, etc:

print(assign_bindings('a, b = (c+d), e'))
print(assign_bindings('a, [b, c, [d, [e]]] = [n1, [n2, n3, [n4, [n5]]]]'))
print(assign_bindings('vals = func(a + b)'))

Output:

[{'a': ['c', 'd'], 'b': ['e']}]
[{'a': ['n1'], 'b': ['n2'], 'c': ['n3'], 'd': ['n4'], 'e': ['n5']}]
[{'vals': ['func', 'a', 'b']}]