How do I extract the names from a simple function?
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
sample spreadsheet http://img441.imageshack.us/img441/1451/84516405.png
**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 so开发者_Python百科me resources here: Python internals: Working with Python ASTs I grabbed a working codegen module here.
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']
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
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').
精彩评论