Deep Black Magic in Python: please help

[Jan Burgy]

Hi everyone,

I am trying to convince my managers that python can replace the outdated and
soon no-longer maintained proprietary system (Tool for Calculator Design) we
use here. In order to achieve this, I need to analyze python code which will
look somethink like this:

def foo(arg_dict):
    return arg_dict["one"] + bar(arg_dict)

def bar(other_dict):
    return other_dict["two"] + other_dict[True and "three" or "four"]

The result of my analysis should return a list

    ['one', 'two', 'three']

Alright, that doesn't sound so bad. With a little RTFM-ing and hacking, I got

-- black magic starts here --

import compiler

def _evaluate(partial, name=""):
    from compiler.ast import Expression
    from compiler.misc import set_filename
    from compiler.pycodegen import ExpressionCodeGenerator

    tree = Expression(partial)
    set_filename(name, tree)
    gen = ExpressionCodeGenerator(tree)
    return eval(gen.getCode())

class RecursiveVisitor(compiler.visitor.ASTVisitor):
    def __init__(self, name):
        self.name = "Name(%s)" % repr(name)
        self.names = {}
    def visitSubscript(self, node, *args):
        if repr(node.expr) == self.name:
            try:
                name = _evaluate(node.subs[0])
            except:
                name = str(node.subs[0])
            self.names[name] = 1
    def visitCallFunc(self, node, *args):
        try:
            from inspect import getsource, getargspec
            from compiler import parse, walk
            func = _evaluate(node.node)
            pos = map(repr, node.args).index(self.name)
            src, arg = getsource(func), getargspec(func)[0]
            tmp, self.name = self.name, "Name(%s)" % repr(arg[pos])
            walk(parse(src), self)
            self.name = tmp
        except Exception, e:
            print str(e)

if __name__ == "__main__":
        
    from inspect import getsource
    from compiler import parse, walk

    src = getsource(foo)
    mod = parse(src)
    visitor = RecursiveVisitor("kw")
    walk(mod, visitor)
    print visitor.names.keys()

-- black magic ends here, ouch --

Once again, I'm in total awe at the possibilities python offers us lowly
users. Unfortunately we're all the same: you give us a finger, we want the
whole arm! I want this method to generalize to method calls as in

class baz:

    def foo(self, arg_dict):
        return arg_dict["one"] + self.bar(arg_dict)

    def bar(self, other_dict):
        return other_dict["two"] + other_dict[True and "three" or "four"]

It shouldn't be all that hard. My problem is the lookup of 'self.bar'. In
the AST it looks something like

    CallFunc(Getattr(Name('self'), 'bar'), [Name('arg_dict')], None, None)

How do I find the corresponding function? Anybody feels like wrapping
their head on this?

Cheers,

Jan Burgy