Source

pycode / typewatch.py

Full commit
import sys
import inspect
import linecache


class Parameters(object):
    
    def __init__(self):
        self.types = None
        self.rtype = None
        
    def __repr__(self):
        return repr(self.__dict__)


class Method(object):
    
    def __init__(self):
        self.callers = []
        self.defined_at = None
        self.parameters = {}

    def __repr__(self):
        return repr(self.__dict__)

    
class Class(object):
    
    def __init__(self):
        self.methods = {}




class Tracer(object):
    

    def __init__(self):
        self.classes = {None: {None:{}}}
        self.current_call = []

        
    def enable(self):
        sys.settrace(self.trace_it)

        
    def disable(self):
        sys.settrace(None)


    def trace_it(self, frame, event, arg):
        # http://docs.python.org/library/inspect.html
        if event == "line":
            lineno = frame.f_lineno
            filename = frame.f_globals["__file__"]
            if filename == "<stdin>":
                filename = "__file__"
            if (filename.endswith(".pyc") or filename.endswith(".pyo")):
                filename = filename[:-1]
            name = frame.f_globals["__name__"]
            line = linecache.getline(filename, lineno)
            #print "LINE %s:%d: %s" % (name, lineno, line.rstrip())
        elif event in ("call", "c_call"):
            arginfo  = inspect.getargvalues(frame)
            if arginfo.args and arginfo.args[0] == "self":
                # assume method
                self.trace_method(frame, arginfo)
            else:
                self.current_call.append(None)
        elif event in ("return", "c_return"):
            self.trace_return(frame, arg)
        elif event in ("exception", "c_exception"):
            (exception, exc_value, exc_traceback) = arg
            print "RAISE", arg
        else:
            print "EVENT is", event
        return self.trace_it

    
    def trace_return(self, frame, result):
        params = self.current_call.pop()
        if params is not None:
            params.rtype = type(result).__name__

    
    def trace_method(self, frame, arginfo):
        line = frame.f_lineno
        arg_names, arg_values = arginfo.args, arginfo.locals
        obj = arginfo.locals["self"]
        cls_name = obj.__class__.__name__
        method_name = frame.f_code.co_name
        arg_types = []
        for arg in arg_names:
            aval = arg_values[arg]
            arg_types.append(type(aval).__name__)
        cls = self.classes.get(cls_name)
        if cls is None:
            cls = self.classes[cls_name] = Class()
        arg_names = tuple(arg_names[1:])
        arg_types = tuple(arg_types[1:])
        pkey = (arg_names, arg_types)
        method = cls.methods.get(method_name)
        if method is None:
            method = cls.methods[method_name] = Method()
        params = method.parameters.get(pkey)
        if params is None:
            params = method.parameters[pkey] = Parameters()
            method.defined_at = line
        method.callers.append(frame.f_back.f_lineno)
        self.current_call.append(params)


def test_func(param):
#    print param
    return param


class TestObj(object):

    def __init__(self):
        pass

    def test_method(self, mp=123):
        return mp * 2


def main():
    foo = "bar"
    #test_func(1)
    o = TestObj()
    o.test_method(4)
    o.test_method(mp=5.6)

t = Tracer()
t.enable()
main()
t.disable()
import pprint
for cls_name, cls in t.classes.items():
    if cls_name is None:
        continue
    print "class", cls_name, cls
    for method_name, method in cls.methods.items():
        for (arg_names, arg_types), param in method.parameters.items():
            print "\tdef", method_name, zip(arg_names, arg_types), "->", param.rtype