root/branches/mk/cheesecake/model.py

"""
-Code borrowed from Michael Hudson's docextractor package with the author's permission.
-The original code is available at <http://codespeak.net/svn/user/mwh/docextractor/>
"""

from compiler import ast
import sys
import os
import cPickle as pickle
import __builtin__
import sets

import compiler
from compiler.transformer import parse, parseFile
from compiler.visitor import walk

class Documentable(object):
    def __init__(self, system, prefix, name, docstring, parent=None):
        self.system = system
        self.prefix = prefix
        self.name = name
        self.docstring = docstring
        self.parent = parent
        self.setup()
    def setup(self):
        self.contents = {}
        self.orderedcontents = []
        self._name2fullname = {}
    def fullName(self):
        return self.prefix + self.name
    def shortdocstring(self):
        docstring = self.docstring
        if docstring:
            docstring = docstring.rstrip()
            if len(docstring) > 20:
                docstring = docstring[:8] + '...' + docstring[-8:]
        return docstring
    def __repr__(self):
        return "%s %r"%(self.__class__.__name__, self.fullName())
    def name2fullname(self, name):
        if name in self._name2fullname:
            return self._name2fullname[name]
        else:
            return self.parent.name2fullname(name)


class Package(Documentable):
    def name2fullname(self, name):
        raise NameError

class Module(Documentable):
    def name2fullname(self, name):
        if name in self._name2fullname:
            return self._name2fullname[name]
        elif name in __builtin__.__dict__:
            return name
        else:
            self.system.warning("optimistic name resolution", name)
            return name

class Class(Documentable):
    def setup(self):
        super(Class, self).setup()
        self.bases = []
    def __repr__(self):
        return "%s(%r, %r) # %r"%(self.__class__.__name__,
                                  self.name, self.shortdocstring(),
                                  self.bases)

class Function(Documentable):
    pass


class System(object):
    Class = Class
    Module = Module
    Package = Package
    Function = Function

    def __init__(self):
        self.current = None
        self._stack = []
        self.allobjects = {}
        self.orderedallobjects = []
        self.rootobjects = []
        self.warnings = {}
        # importstargraph contains edges {importedby:[imports]} but only
        # for import * statements
        self.importstargraph = {}
        self.func_called = {}

    def _push(self, cls, name, docstring):
        if self.current:
            prefix = self.current.fullName() + '.'
            parent = self.current
        else:
            prefix = ''
            parent = None
        obj = cls(self, prefix, name, docstring, parent)
        if parent:
            parent.orderedcontents.append(obj)
            parent.contents[name] = obj
            parent._name2fullname[name] = obj.fullName()
        else:
            self.rootobjects.append(obj)
        self.current = obj
        self.orderedallobjects.append(obj)
        fullName = obj.fullName()
        #print 'push', cls.__name__, fullName
        if fullName in self.allobjects:
            obj = self.handleDuplicate(obj)
        else:
            self.allobjects[obj.fullName()] = obj
        return obj

    def handleDuplicate(self, obj):
        '''This is called when we see two objects with the same
        .fullName(), for example:

        class C:
            if something:
                def meth(self):
                    implementation 1
            else:
                def meth(self):
                    implementation 2

        The default is that the second definition "wins".
        '''
        self.warning("duplicate", self.allobjects[obj.fullName()])
        self.allobjects[obj.fullName()] = obj
        return obj


    def _pop(self, cls):
        assert isinstance(self.current, cls)
##         if self.current.parent:
##             print 'pop', self.current.fullName(), '->', self.current.parent.fullName()
##         else:
##             print 'pop', self.current.fullName(), '->', self.current.parent
        self.current = self.current.parent

    def push(self, obj, node=None):
        self._stack.append(self.current)
        self.current = obj

    def pop(self, obj):
        assert self.current is obj, "%r is not %r"%(self.current, obj)
        self.current = self._stack.pop()

    def pushClass(self, name, docstring):
        return self._push(self.Class, name, docstring)
    def popClass(self):
        self._pop(self.Class)

    def pushModule(self, name, docstring):
        return self._push(self.Module, name, docstring)
    def popModule(self):
        self._pop(self.Module)

    def pushFunction(self, name, docstring, func_called):
        self.func_called.update(func_called)
        return self._push(self.Function, name, docstring)
    def popFunction(self):
        self._pop(self.Function)

    def pushPackage(self, name, docstring):
        return self._push(self.Package, name, docstring)
    def popPackage(self):
        self._pop(self.Package)

    def report(self):
        for o in self.rootobjects:
            self._report(o, '')

    def _report(self, o, indent):
        print indent, o
        for o2 in o.orderedcontents:
            self._report(o2, indent+'  ')

    def resolveAlias(self, n):
        if '.' not in n:
            return n
        mod, clsname = n.split('.')
        if not mod or mod not in self.allobjects:
            return n
        m = self.allobjects[mod]
        if not isinstance(m, Module):
            return n
        if clsname in m._name2fullname:
            newname = m.name2fullname(clsname)
            if newname not in self.allobjects:
                return self.resolveAlias(newname)
            else:
                return newname

    def resolveAliases(self):
        for ob in self.orderedallobjects:
            if not isinstance(ob, Class):
                continue
            for i, b in enumerate(ob.bases):
                if b not in self.allobjects:
                    ob.bases[i] = self.resolveAlias(b)

    def warning(self, type, detail):
        fn = self.current.fullName()
        #print fn, type, detail
        self.warnings.setdefault(type, []).append((fn, detail))

    def objectsOfType(self, cls):
        for o in self.orderedallobjects:
            if isinstance(o, cls):
                yield o

def expandModname(system, modname, givewarning=True):
    c = system.current
    if '.' in modname:
        prefix, suffix = modname.split('.', 1)
        suffix = '.' + suffix
    else:
        prefix, suffix = modname, ''
    while c is not None and not isinstance(c, Package):
        c = c.parent
    while c is not None:
        if prefix in c.contents:
            break
        c = c.parent
    if c is not None:
        if givewarning:
            system.warning("local import", modname)
        return c.contents[prefix].fullName() + suffix
    else:
        return prefix + suffix

class ImportStarFinder(object):
    def __init__(self, system, modfullname):
        self.system = system
        self.modfullname = modfullname

    def visitFrom(self, node):
        if node.names[0][0] == '*':
            modname = expandModname(self.system, node.modname, False)
            self.system.importstargraph.setdefault(
                modname, []).append(self.modfullname)

class ModuleVistor(object):
    def __init__(self, system, modname):
        self.system = system
        self.modname = modname
        self.morenodes = []

    def default(self, node):
        for child in node.getChildNodes():
            self.visit(child)

    def postpone(self, docable, node):
        self.morenodes.append((docable, node))

    def visitModule(self, node):
        if self.system.current and self.modname in self.system.current.contents:
            m = self.system.current.contents[self.modname]
            assert m.docstring is None
            m.docstring = node.doc
            self.system.push(m, node)
            self.default(node)
            self.system.pop(m)
        else:
            self.system.pushModule(self.modname, node.doc)
            self.default(node)
            self.system.popModule()

    def visitClass(self, node):
        cls = self.system.pushClass(node.name, node.doc)
        for n in node.bases:
            if isinstance(n, ast.Name):
                cls.bases.append(cls.parent.name2fullname(n.name))
            elif isinstance(n, ast.Getattr):
                p = []
                while isinstance(n, ast.Getattr):
                    p.append(n.attrname)
                    n = n.expr
                assert isinstance(n, ast.Name)
                p.append(cls.parent.name2fullname(n.name))
                p.reverse()
                assert None not in p, n
                cls.bases.append('.'.join(p))
            else:
                assert not n
        self.default(node)
        self.system.popClass()

    def visitFrom(self, node):
        modname = expandModname(self.system, node.modname)
        name2fullname = self.system.current._name2fullname
        for fromname, asname in node.names:
            if fromname == '*':
                self.system.warning("import *", modname)
                if modname not in self.system.allobjects:
                    return
                mod = self.system.allobjects[modname]
                #snarl (see below)
                #assert mod.processed
                self.system.warning("mwh is an idiot", "")
                for n in mod.contents:
                    name2fullname[n] = modname + '.' + n
                return
            if asname is None:
                asname = fromname
            name2fullname[asname] = modname + '.' + fromname

    def visitImport(self, node):
        name2fullname = self.system.current._name2fullname
        for fromname, asname in node.names:
            fullname = expandModname(self.system, fromname)
            if asname is None:
                asname = fromname.split('.', 1)[0]
                # aaaaargh! python sucks.
                parts = fullname.split('.')
                for i, part in enumerate(fullname.split('.')[::-1]):
                    if part == asname:
                        fullname = '.'.join(parts[:len(parts)-i])
                        name2fullname[asname] = fullname
                        break
                else:
                    name2fullname[asname] = '.'.join(parts)
            else:
                name2fullname[asname] = fullname


    def visitFunction(self, node):
        fc = {}
        get_function_calls(node, fc)
        #print fc.keys()
        func = self.system.pushFunction(node.name, node.doc, fc)
        # ast.Function has a pretty lame representation of
        # arguments. Let's convert it to a nice concise
        # getargspec-like format and include it in the Function
        # object.
        argnames = node.argnames[:]
        kwname = starargname = None
        if node.kwargs:
            kwname = argnames.pop(-1)
        if node.varargs:
            starargname = argnames.pop(-1)
        defaults = []
        for default in node.defaults:
            if isinstance(default, ast.Const):
                defaults.append(default.value)
            elif isinstance(default, ast.Name):
                defaults.append(default.name)
            else:
                self.system.warning("unparseable default", repr(default))
                defaults.append('???')
                #assert False, "don't know how to handle default %r"%(default,)
        # argh, convert unpacked-arguments from tuples to lists,
        # because that's what getargspec uses and the unit test
        # compares it
        argnames2 = []
        for argname in argnames:
            if isinstance(argname, tuple):
                argname = list(argname)
            argnames2.append(argname)
        func.argspec = (argnames2, starargname, kwname, tuple(defaults))
        #for child in node.getChildren():
        #    if isinstance(child, compiler.ast.Stmt):
        #        for c in child.getChildren():
        #            print c.__class__
        #            print c
        self.postpone(func, node.code)
        self.system.popFunction()

def get_function_calls(node, fc):
    if not isinstance(node, compiler.ast.Node):
        return
    for child in node.getChildren():
        #print "child:", child
        if isinstance(child, compiler.ast.CallFunc):
                funcname = ""
                attrname = ""
                n = child.node
                #print "n:", n
                #print n.__class__
                if isinstance(n, compiler.ast.Getattr):
                    expr = n.expr
                    if isinstance(expr, compiler.ast.Name):
                        funcname = expr.name
                    attrname = n.attrname
                func_called = ""
                if funcname: func_called = funcname + "."
                func_called += attrname
                if func_called:
                    fc[func_called] = 1
        get_function_calls(child, fc)
    
def processModuleAst(ast, name, system):
    mv = ModuleVistor(system, name)
    walk(ast, mv)
    while mv.morenodes:
        obj, node = mv.morenodes.pop(0)
        system.push(obj, node)
        mv.visit(node)
        system.pop(obj)


def fromText(src, modname='<test>', system=None):
    if system is None:
        system = System()
    processModuleAst(parse(src), modname, system)
    return system.rootobjects[0]


def preprocessDirectory(system, dirpath):
    package = system.pushPackage(os.path.basename(dirpath), None)
    for fname in os.listdir(dirpath):
        fullname = os.path.join(dirpath, fname)
        if os.path.isdir(fullname) and os.path.exists(os.path.join(fullname, '__init__.py')):
            preprocessDirectory(system, fullname)
        elif fname.endswith('.py'):
            modname = os.path.splitext(fname)[0]
            mod = system.pushModule(modname, None)
            mod.filepath = fullname
            mod.processed = False
            system.popModule()
    system.popPackage()

def processDirectory(system, dirpath):
    preprocessDirectory(system, dirpath)
    modlist = list(system.objectsOfType(Module))
    for mod in modlist:
        system.push(mod.parent)
        isf = ImportStarFinder(system, mod.fullName())
        walk(parseFile(mod.filepath), isf)
        system.pop(mod.parent)

    # snarl; a toposort is meant to go here.
    newlist = modlist

    for mod in newlist:
        system.push(mod.parent)
        processModuleAst(parseFile(mod.filepath), mod.name, system)
        mod.processed = True
        system.pop(mod.parent)

def main(argv):
    if '-r' in argv:
        argv.remove('-r')
        assert len(argv) == 1
        system = System()
        processDirectory(system, argv[0])
        pickle.dump(system, open('da.out', 'wb'), pickle.HIGHEST_PROTOCOL)
        print
        print 'warning summary:'
        for k, v in system.warnings.iteritems():
            print k, len(v)
    else:
        system = System()
        for fname in argv:
            modname = os.path.splitext(os.path.basename(fname))[0] # XXX!
            processModuleAst(parseFile(fname), modname, system)
        system.report()



if __name__ == '__main__':
    main(sys.argv[1:])