Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

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 some resources here: Python internals: Working with Python ASTs I grabbed a working codegen module here.

like image 536
Eddy Pronk Avatar asked Jul 09 '10 13:07

Eddy Pronk


1 Answers

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']
like image 192
Eddy Pronk Avatar answered Sep 19 '22 14:09

Eddy Pronk