Python AST processing

Question

I have a Python AST [as returned by ast.parse()].

I know this is an AST of a class method.

How do I find all calls to other methods of the same class?

Basically, I want to collect something like:

['foo', 'bar']

for a code snippet like:

def baz(self): # this is a class method
    '''baz docstring'''
    self.foo() + self.bar()

I need a function that will accept an AST and will return the list of other methods [method names as strings] of the class that are being invoked inside a method of the same class.

Answer 1

The general approach is to subclass ast.NodeVisitor:

>>> class VisitCalls(ast.NodeVisitor):
...   def visit_Call(self, what):
...     if what.func.value.id == 'self':
...       print what.func.attr
... 
>>> f='''def x(self):
...   return self.bar() + self.baz()
... '''
>>> xx = ast.parse(f)
>>> VisitCalls().visit(xx)
bar
baz

However, this will only catch "immediate" calls to self.something. In the general case you could have e.g. somelist.append(self.blah) and then much later in the code somelist[i + j](): the problem of determining whether the latter is a call to self.blah or to some other callable that has nothing to do with methods of the current instance is Turing-complete (CS jargon for "completely insoluble in the general case", much like a mathematician might say "NP-hard";-).

But if all you need is to solve the simple "immediate call" case, you're good to go;-).