/py/_code/source.py

from __future__ import generators
import sys
import inspect, tokenize
import py
from types import ModuleType
cpy_compile = compile

try:
    import _ast
    from _ast import PyCF_ONLY_AST as _AST_FLAG
except ImportError:
    _AST_FLAG = 0
    _ast = None


class Source(object):
    """ a immutable object holding a source code fragment,
        possibly deindenting it.
    """
    _compilecounter = 0
    def __init__(self, *parts, **kwargs):
        self.lines = lines = []
        de = kwargs.get('deindent', True)
        rstrip = kwargs.get('rstrip', True)
        for part in parts:
            if not part:
                partlines = []
            if isinstance(part, Source):
                partlines = part.lines
            elif isinstance(part, (tuple, list)):
                partlines = [x.rstrip("\n") for x in part]
            elif isinstance(part, py.builtin._basestring):
                partlines = part.split('\n')
                if rstrip:
                    while partlines:
                        if partlines[-1].strip():
                            break
                        partlines.pop()
            else:
                partlines = getsource(part, deindent=de).lines
            if de:
                partlines = deindent(partlines)
            lines.extend(partlines)

    def __eq__(self, other):
        try:
            return self.lines == other.lines
        except AttributeError:
            if isinstance(other, str):
                return str(self) == other
            return False

    def __getitem__(self, key):
        if isinstance(key, int):
            return self.lines[key]
        else:
            if key.step not in (None, 1):
                raise IndexError("cannot slice a Source with a step")
            return self.__getslice__(key.start, key.stop)

    def __len__(self):
        return len(self.lines)

    def __getslice__(self, start, end):
        newsource = Source()
        newsource.lines = self.lines[start:end]
        return newsource

    def strip(self):
        """ return new source object with trailing
            and leading blank lines removed.
        """
        start, end = 0, len(self)
        while start < end and not self.lines[start].strip():
            start += 1
        while end > start and not self.lines[end-1].strip():
            end -= 1
        source = Source()
        source.lines[:] = self.lines[start:end]
        return source

    def putaround(self, before='', after='', indent=' ' * 4):
        """ return a copy of the source object with
            'before' and 'after' wrapped around it.
        """
        before = Source(before)
        after = Source(after)
        newsource = Source()
        lines = [ (indent + line) for line in self.lines]
        newsource.lines = before.lines + lines +  after.lines
        return newsource

    def indent(self, indent=' ' * 4):
        """ return a copy of the source object with
            all lines indented by the given indent-string.
        """
        newsource = Source()
        newsource.lines = [(indent+line) for line in self.lines]
        return newsource

    def getstatement(self, lineno, assertion=False):
        """ return Source statement which contains the
            given linenumber (counted from 0).
        """
        start, end = self.getstatementrange(lineno, assertion)
        return self[start:end]

    def getstatementrange(self, lineno, assertion=False):
        """ return (start, end) tuple which spans the minimal
            statement region which containing the given lineno.
            raise an IndexError if no such statementrange can be found.
        """
        # XXX there must be a better than these heuristic ways ...
        # XXX there may even be better heuristics :-)
        if not (0 <= lineno < len(self)):
            raise IndexError("lineno out of range")

        # 1. find the start of the statement
        from codeop import compile_command
        end = None
        for start in range(lineno, -1, -1):
            if assertion:
                line = self.lines[start]
                # the following lines are not fully tested, change with care
                if 'super' in line and 'self' in line and '__init__' in line:
                    raise IndexError("likely a subclass")
                if "assert" not in line and "raise" not in line:
                    continue
            trylines = self.lines[start:lineno+1]
            # quick hack to indent the source and get it as a string in one go
            trylines.insert(0, 'if xxx:')
            trysource = '\n '.join(trylines)
            #              ^ space here
            try:
                compile_command(trysource)
            except (SyntaxError, OverflowError, ValueError):
                continue

            # 2. find the end of the statement
            for end in range(lineno+1, len(self)+1):
                trysource = self[start:end]
                if trysource.isparseable():
                    return start, end
        if end is None:
            raise IndexError("no valid source range around line %d " % (lineno,))
        return start, end

    def getblockend(self, lineno):
        # XXX
        lines = [x + '\n' for x in self.lines[lineno:]]
        blocklines = inspect.getblock(lines)
        #print blocklines
        return lineno + len(blocklines) - 1

    def deindent(self, offset=None):
        """ return a new source object deindented by offset.
            If offset is None then guess an indentation offset from
            the first non-blank line.  Subsequent lines which have a
            lower indentation offset will be copied verbatim as
            they are assumed to be part of multilines.
        """
        # XXX maybe use the tokenizer to properly handle multiline
        #     strings etc.pp?
        newsource = Source()
        newsource.lines[:] = deindent(self.lines, offset)
        return newsource

    def isparseable(self, deindent=True):
        """ return True if source is parseable, heuristically
            deindenting it by default.
        """
        try:
            import parser
        except ImportError:
            syntax_checker = lambda x: compile(x, 'asd', 'exec')
        else:
            syntax_checker = parser.suite

        if deindent:
            source = str(self.deindent())
        else:
            source = str(self)
        try:
            #compile(source+'\n', "x", "exec")
            syntax_checker(source+'\n')
        except KeyboardInterrupt:
            raise
        except Exception:
            return False
        else:
            return True

    def __str__(self):
        return "\n".join(self.lines)

    def compile(self, filename=None, mode='exec',
                flag=generators.compiler_flag,
                dont_inherit=0, _genframe=None):
        """ return compiled code object. if filename is None
            invent an artificial filename which displays
            the source/line position of the caller frame.
        """
        if not filename or py.path.local(filename).check(file=0):
            if _genframe is None:
                _genframe = sys._getframe(1) # the caller
            fn,lineno = _genframe.f_code.co_filename, _genframe.f_lineno
            base = "<%d-codegen " % self._compilecounter
            self.__class__._compilecounter += 1
            if not filename:
                filename = base + '%s:%d>' % (fn, lineno)
            else:
                filename = base + '%r %s:%d>' % (filename, fn, lineno)
        source = "\n".join(self.lines) + '\n'
        try:
            co = cpy_compile(source, filename, mode, flag)
        except SyntaxError:
            ex = sys.exc_info()[1]
            # re-represent syntax errors from parsing python strings
            msglines = self.lines[:ex.lineno]
            if ex.offset:
                msglines.append(" "*ex.offset + '^')
            msglines.append("(code was compiled probably from here: %s)" % filename)
            newex = SyntaxError('\n'.join(msglines))
            newex.offset = ex.offset
            newex.lineno = ex.lineno
            newex.text = ex.text
            raise newex
        else:
            if flag & _AST_FLAG:
                return co
            lines = [(x + "\n") for x in self.lines]
            if sys.version_info[0] >= 3:
                # XXX py3's inspect.getsourcefile() checks for a module
                # and a pep302 __loader__ ... we don't have a module
                # at code compile-time so we need to fake it here
                m = ModuleType("_pycodecompile_pseudo_module")
                py.std.inspect.modulesbyfile[filename] = None
                py.std.sys.modules[None] = m
                m.__loader__ = 1
            py.std.linecache.cache[filename] = (1, None, lines, filename)
            return co

#
# public API shortcut functions
#

def compile_(source, filename=None, mode='exec', flags=
            generators.compiler_flag, dont_inherit=0):
    """ compile the given source to a raw code object,
        and maintain an internal cache which allows later
        retrieval of the source code for the code object
        and any recursively created code objects.
    """
    if _ast is not None and isinstance(source, _ast.AST):
        # XXX should Source support having AST?
        return cpy_compile(source, filename, mode, flags, dont_inherit)
    _genframe = sys._getframe(1) # the caller
    s = Source(source)
    co = s.compile(filename, mode, flags, _genframe=_genframe)
    return co


def getfslineno(obj):
    """ Return source location (path, lineno) for the given object.
    If the source cannot be determined return ("", -1)
    """
    try:
        code = py.code.Code(obj)
    except TypeError:
        try:
            fn = (py.std.inspect.getsourcefile(obj) or
                  py.std.inspect.getfile(obj))
        except TypeError:
            return "", -1

        fspath = fn and py.path.local(fn) or None
        lineno = -1
        if fspath:
            try:
                _, lineno = findsource(obj)
            except IOError:
                pass
    else:
        fspath = code.path
        lineno = code.firstlineno
    assert isinstance(lineno, int)
    return fspath, lineno

#
# helper functions
#

def findsource(obj):
    try:
        sourcelines, lineno = py.std.inspect.findsource(obj)
    except py.builtin._sysex:
        raise
    except:
        return None, -1
    source = Source()
    source.lines = [line.rstrip() for line in sourcelines]
    return source, lineno

def getsource(obj, **kwargs):
    obj = py.code.getrawcode(obj)
    try:
        strsrc = inspect.getsource(obj)
    except IndentationError:
        strsrc = "\"Buggy python version consider upgrading, cannot get source\""
    assert isinstance(strsrc, str)
    return Source(strsrc, **kwargs)

def deindent(lines, offset=None):
    if offset is None:
        for line in lines:
            line = line.expandtabs()
            s = line.lstrip()
            if s:
                offset = len(line)-len(s)
                break
        else:
            offset = 0
    if offset == 0:
        return list(lines)
    newlines = []
    def readline_generator(lines):
        for line in lines:
            yield line + '\n'
        while True:
            yield ''

    r = readline_generator(lines)
    try:
        readline = r.next
    except AttributeError:
        readline = r.__next__

    try:
        for _, _, (sline, _), (eline, _), _ in tokenize.generate_tokens(readline):
            if sline > len(lines):
                break # End of input reached
            if sline > len(newlines):
                line = lines[sline - 1].expandtabs()
                if line.lstrip() and line[:offset].isspace():
                    line = line[offset:] # Deindent
                newlines.append(line)

            for i in range(sline, eline):
                # Don't deindent continuing lines of
                # multiline tokens (i.e. multiline strings)
                newlines.append(lines[i])
    except (IndentationError, tokenize.TokenError):
        pass
    # Add any lines we didn't see. E.g. if an exception was raised.
    newlines.extend(lines[len(newlines):])
    return newlines