How can I retrieve function names and attributes from Python code with AST?

user_8450935 0
Viewed 295

Given a piece of code like this:

aaa = bbb()
ccc.ddd().eee(fff(foo))
ggg(bar)
hhh().iii

I would like to get something like this:

['bbb', 'ccc.ddd.eee', 'fff', 'ggg', 'hhh.iii']

I'm aware of the visit_Call and visit_Name methods, but I can't get them to return only function names or function names with functions called on objects returned by those functions (like foo().bar()).

abstract-syntax-tree python
1 Answers

To produce function-attribute lookup paths, you will want to traverse the ast.AST object, checking for any ast.Calls or ast.Attributes, and if any are found, you will then need to walk those objects, saving the names and attributes as they occur. The recursive pattern below traverses the main AST (parse) and utilizes a helper function (parse_chain) to walk any attributes or calls for names present:

import ast
def parse(d, c):
  def parse_chain(d, c, p=[]):
     if isinstance(d, ast.Name):
        return [d.id]+p
     if isinstance(d, ast.Call):
        for i in d.args:
           parse(i, c)
        return parse_chain(d.func, c, p)
     if isinstance(d, ast.Attribute):
        return parse_chain(d.value, c, [d.attr]+p)
  if isinstance(d, (ast.Call, ast.Attribute)):
     c.append('.'.join(parse_chain(d, c)))
  else:
     for i in getattr(d, '_fields', []):
       if isinstance(t:=getattr(d, i), list):
          for i in t:
             parse(i, c)
       else:
          parse(t, c)

results = []
s = """
aaa = bbb()
ccc.ddd().eee(fff(foo))
ggg(bar)
hhh().iii
"""
parse(ast.parse(s), results)
print(results)

Output:

['bbb', 'fff', 'ccc.ddd.eee', 'ggg', 'hhh.iii']
user_73267380
Related