tonfa/sphinx / sphinx / pycode /

# -*- coding: utf-8 -*-

    Utilities parsing and analyzing Python code.

    :copyright: Copyright 2007-2009 by the Sphinx team, see AUTHORS.
    :license: BSD, see LICENSE for details.

import sys
from os import path
from cStringIO import StringIO

from sphinx.pycode import nodes
from sphinx.pycode.pgen2 import driver, token, tokenize, parse, literals
from sphinx.util.docstrings import prepare_docstring, prepare_commentdoc

# load the Python grammar
_grammarfile = path.join(path.dirname(__file__), 'Grammar.txt')
pygrammar = driver.load_grammar(_grammarfile)
pydriver = driver.Driver(pygrammar, convert=nodes.convert)

# an object with attributes corresponding to token and symbol names
class sym: pass
for k, v in pygrammar.symbol2number.iteritems():
    setattr(sym, k, v)
for k, v in token.tok_name.iteritems():
    setattr(sym, v, k)

# a dict mapping terminal and nonterminal numbers to their names
number2name = pygrammar.number2symbol.copy()

_eq = nodes.Leaf(token.EQUAL, '=')

class AttrDocVisitor(nodes.NodeVisitor):
    Visitor that collects docstrings for attribute assignments on toplevel and
    in classes.

    The docstrings can either be in special '#:' comments before the assignment
    or in a docstring after it.
    def init(self, scope):
        self.scope = scope
        self.namespace = []
        self.collected = {}

    def visit_classdef(self, node):

    def visit_expr_stmt(self, node):
        """Visit an assignment which may have a special comment before it."""
        if _eq not in node.children:
            # not an assignment (we don't care for augmented assignments)
        pnode = node[0]
        prefix = pnode.get_prefix()
        # if the assignment is the first statement on a new indentation
        # level, its preceding whitespace and comments are not assigned
        # to that token, but the first INDENT or DEDENT token
        while not prefix:
            pnode = pnode.get_prev_leaf()
            if not pnode or pnode.type not in (token.INDENT, token.DEDENT):
            prefix = pnode.get_prefix()
        docstring = prepare_commentdoc(prefix)
        if docstring:
            self.add_docstring(node, docstring)

    def visit_simple_stmt(self, node):
        """Visit a docstring statement which may have an assignment before."""
        if node[0].type != token.STRING:
            # not a docstring; but still need to visit children
            return self.generic_visit(node)
        prev = node.get_prev_sibling()
        if not prev:
        if prev.type == sym.simple_stmt and \
               prev[0].type == sym.expr_stmt and _eq in prev[0].children:
            # need to "eval" the string because it's returned in its original form
            docstring = prepare_docstring(literals.evalString(node[0].value))
            self.add_docstring(prev[0], docstring)

    def visit_funcdef(self, node):
        # don't descend into functions -- nothing interesting there

    def add_docstring(self, node, docstring):
        # add an item for each assignment target
        for i in range(0, len(node) - 1, 2):
            target = node[i]
            if target.type != token.NAME:
                # don't care about complex targets
            namespace = '.'.join(self.namespace)
            if namespace.startswith(self.scope):
                self.collected[namespace, target.value] = docstring

class PycodeError(Exception):
    def __str__(self):
        res = self.args[0]
        if len(self.args) > 1:
            res += ' (exception was: %r)' % self.args[1]
        return res

class ModuleAnalyzer(object):
    # cache for analyzer objects -- caches both by module and file name
    cache = {}

    def for_string(cls, string, modname, srcname='<string>'):
        return cls(StringIO(string), modname, srcname)

    def for_file(cls, filename, modname):
        if ('file', filename) in cls.cache:
            return cls.cache['file', filename]
            fileobj = open(filename, 'r')
        except Exception, err:
            raise PycodeError('error opening %r' % filename, err)
        obj = cls(fileobj, modname, filename)
        cls.cache['file', filename] = obj
        return obj

    def for_module(cls, modname):
        if ('module', modname) in cls.cache:
            return cls.cache['module', modname]
        if modname not in sys.modules:
            except ImportError, err:
                raise PycodeError('error importing %r' % modname, err)
        mod = sys.modules[modname]
        if hasattr(mod, '__loader__'):
                source = mod.__loader__.get_source(modname)
            except Exception, err:
                raise PycodeError('error getting source for %r' % modname, err)
            obj = cls.for_string(source, modname)
            cls.cache['module', modname] = obj
            return obj
        filename = getattr(mod, '__file__', None)
        if filename is None:
            raise PycodeError('no source found for module %r' % modname)
        filename = path.normpath(filename)
        lfilename = filename.lower()
        if lfilename.endswith('.pyo') or lfilename.endswith('.pyc'):
            filename = filename[:-1]
        elif not lfilename.endswith('.py'):
            raise PycodeError('source is not a .py file: %r' % filename)
        if not path.isfile(filename):
            raise PycodeError('source file is not present: %r' % filename)
        obj = cls.for_file(filename, modname)
        cls.cache['module', modname] = obj
        return obj

    def __init__(self, source, modname, srcname):
        self.modname = modname
        self.srcname = srcname
        # file-like object yielding source lines
        self.source = source

        # will be filled by tokenize()
        self.tokens = None
        # will be filled by parse()
        self.parsetree = None

    def tokenize(self):
        """Generate tokens from the source."""
        if self.tokens is not None:
        self.tokens = list(tokenize.generate_tokens(self.source.readline))

    def parse(self):
        """Parse the generated source tokens."""
        if self.parsetree is not None:
        self.parsetree = pydriver.parse_tokens(self.tokens)

    def find_attr_docs(self, scope=''):
        """Find class and module-level attributes and their documentation."""
        attr_visitor = AttrDocVisitor(number2name, scope)
        return attr_visitor.collected

    def find_tags(self):
        """Find class, function and method definitions and their location."""
        result = {}
        namespace = []
        stack = []
        indent = 0
        defline = False
        expect_indent = False
        def tokeniter(ignore = (token.COMMENT, token.NL)):
            for tokentup in self.tokens:
                if tokentup[0] not in ignore:
                    yield tokentup
        tokeniter = tokeniter()
        for type, tok, spos, epos, line in tokeniter:
            if expect_indent:
                if type != token.INDENT:
                    # no suite -- one-line definition
                    assert stack
                    dtype, fullname, startline, _ = stack.pop()
                    endline = epos[0]
                    result[fullname] = (dtype, startline, endline)
                expect_indent = False
            if tok in ('def', 'class'):
                name =[1]
                fullname = '.'.join(namespace)
                stack.append((tok, fullname, spos[0], indent))
                defline = True
            elif type == token.INDENT:
                expect_indent = False
                indent += 1
            elif type == token.DEDENT:
                indent -= 1
                # if the stacklevel is the same as it was before the last
                # def/class block, this dedent closes that block
                if stack and indent == stack[-1][3]:
                    dtype, fullname, startline, _ = stack.pop()
                    endline = spos[0]
                    result[fullname] = (dtype, startline, endline)
            elif type == token.NEWLINE:
                # if this line contained a definition, expect an INDENT to start the
                # suite; if there is no such INDENT it's a one-line definition
                if defline:
                    defline = False
                    expect_indent = True
        return result

if __name__ == '__main__':
    import time, pprint
    x0 = time.time()
    #ma = ModuleAnalyzer.for_file(__file__.rstrip('c'), '')
    ma = ModuleAnalyzer.for_file('sphinx/builders/', '')
    x1 = time.time()
    x2 = time.time()
    #for (ns, name), doc in ma.find_attr_docs().iteritems():
    #    print '>>', ns, name
    #    print '\n'.join(doc)
    x3 = time.time()
    #print nodes.nice_repr(ma.parsetree, number2name)
    print "tokenizing %.4f, parsing %.4f, finding %.4f" % (x1-x0, x2-x1, x3-x2)