How do I extract the names from a simple function?

Question

I've got this piece of code:

import inspect
import ast

def func(foo):
    return foo.bar - foo.baz

s = inspect.getsource(func)
xx = ast.parse(s)

class VisitCalls(ast.NodeVisitor):
    def visit_Name(self, what):
        if what.id == 'foo':
            print ast.dump(what.ctx)

VisitCalls().visit(xx)

From function 'func' I'd like to extract:

['foo.bar', 'foo.baz']

or something like:

(('foo', 'bar'), ('foo', 'baz))

edited

Some background to explain why I think I need to do this

I want to convert the code of a trivial python function to a spreadsheet formula.

So I need to convert:

foo.bar - foo.baz

to:

=A1-B1

*edited again

What I've got so far.

The program below outputs:

('A1', 5)
('B1', 3)
('C1', '= A1 - B1')

The code:

import ast, inspect
import codegen # by Armin Ronacher
from collections import OrderedDict

class SpreadSheetFormulaTransformer(ast.NodeTransformer):
    def __init__(self, sym):
        self.sym = sym
    def visit_Attribute(self, node):
        name = self.sym[id(eval(codegen.to_source(node)))]
        return ast.Name(id=name, ctx=ast.Load())

def create(**kwargs):
    class Foo(object): pass
    x = Foo()
    x.__dict__.update(kwargs)
    return x

def register(x,y):
    cell[y] = x
    sym[id(x)] = y

def func(foo):
    return foo.bar - foo.baz

foo = create(bar=5, baz=3)
cell = OrderedDict()
sym = {}

register(foo.bar, 'A1')
register(foo.baz, 'B1')

source = inspect.getsource(func)
tree = ast.parse(source)
guts = tree.body[0].body[0].value
SpreadSheetFormulaTransformer(sym).visit(guts)

code = '= ' + codegen.to_source(guts)
cell['C1'] = code

for x in cell.iteritems():
    print x

I found some resources here: Python internals: Working with Python ASTs I grabbed a working codegen module here.

Answer 1

I am not sure why you need to retirieve names, a very crude way to get all names and dots in function is

import inspect
import parser
import symbol
import token
import pprint

def func(foo):
    return foo.bar - foo.baz

s = inspect.getsource(func)
st = parser.suite(s)

def search(st):
    if not isinstance(st, list):
        return
    if st[0] in [token.NAME, token.DOT]:
        print st[1],
    else:
        for s in st[1:]:
            search(s)

search(parser.ast2list(st))

output:

def func foo return foo . bar foo . baz

May be you can improve upon that by reading syntax tree more elegantly, I am using parser instead of ast module because i am on python 2.5

Answer 2

I haven't used the new ast module yet, but I've working code that uses the older compiler.ast to achieve something similar:

    def visitGetattr(self, node):
        full_name = [node.attrname]
        parent = node.expr
        while isinstance(parent, compiler.ast.Getattr):
            full_name.append(parent.attrname)
            parent = parent.expr
        if isinstance(parent, compiler.ast.Name):
            full_name.append(parent.name)
            full_name = ".".join(reversed(full_name))
            # do something with full_name
        for c in node.getChildNodes():
            self.visit(c)

Code slightly paraphrased, I may have introduced inadvertent bugs. I hope this gives you the general idea: you need to visit both Name and Getattr nodes and construct dotted names, and also deal with the fact that you'll see all the intermediate values too (e.g. 'foo' and 'foo.bar').

Answer 3

import ast, inspect
import codegen # by Armin Ronacher

def func(foo):
    return foo.bar - foo.baz

names = []

class CollectAttributes(ast.NodeVisitor):
    def visit_Attribute(self, node):
        names.append(codegen.to_source(node))

source = inspect.getsource(func)

tree = ast.parse(source)
guts = tree.body[0].body[0].value
CollectAttributes().visit(guts)
print names

output:

['foo.bar', 'foo.baz']