Untitled

import inspect
import code
import re
import editor
import tokenize
import linecache
import shutil
import time
import functools


global tweak_history, tweak_originals
tweak_history = {} # Stores a list of versions of the source code of a function
tweak_originals = {} # Stores a

#Only wrapped in a class so that I can just use T() instead of T.tweak() while still having access to T.apply() etc.

class Tweaker(object):
    @staticmethod
    def __call__(ob):
        """Same as tweak(ob)"""
        Tweaker.tweak(ob)

    @staticmethod
    def tweak(ob):
        """
            Edit the source code of a function using an external
            editor, and then - without reloading - use the new
            function definition in your current code.

            If you liked the tweaks you made, you can permenantly
            apply() them, or if they broke stuff, revert() them.

            Methods of instances use the same function as their
            class, so changing the method of one instance changes
            them ALL.

            The function is looked up from this method's argument
            using getfunction() so see notes there for what to pass.

            Be careful while editing, do not play around with the
            indentation too much, you are liable to break things.
        """

        olds = getsource(ob)
        news = editor.edit(olds, filesuffix=".py")
        if news:
            setsource(ob, news)
        else:
            print "Nothing changed"

    @staticmethod
    def getfunction(orig):
        """
            Get the actual function from the argument - tries the following:

            If we are passed a string, eval it in all the stack frames above us until it works
            If we have a partial, get it's .func property
            If we have a method, get it's im_func property
            If we have a function that looks like a decorator.decorator decorated one, get the original function
            If we have a function with a name that doesn't match the string input's name, try and jump out of a closure
        """
        output = orig
        name = None

        #If we are passed a string, first try and find what it refers to,
        # We will use the name to attempt to jump out of a decorator closure.
        if type(orig) is str:
            for stackframe in [x[0] for x in inspect.stack()[2:]]:
                try:
                    output = eval(orig, stackframe.f_locals, stackframe.f_globals)
                    try:
                        name = orig[orig.rindex(".")+1:]
                    except:
                        name = orig
                    break
                except:
                    pass

        #Get the original function that a partial refers to
        if type(output) is functools.partial:
            output = output.func

        #Get the actual function a method refers to
        if inspect.ismethod(orig):
            output = output.im_func

        #This is the only way I've found to identify functions decorated by decorator.decorator
        if inspect.isfunction(output) and '_func_' in output.func_globals:
            output = output.func_globals['_func_']

        #What about callables masquerading as functions?
        if callable(output) and not inspect.isfunction(output):
            output = output.__call__

        #Now we try and guess our way out of a normal decorator closure
        if name and inspect.isfunction(output) and output.__name__ != name:
            for idx in range(len(output.func_closure)):
                poss = output.func_closure[idx].cell_contents
                if inspect.isfunction(poss) and poss.__name__ == name:
                    output = poss

        if not inspect.isfunction(output):
            raise Exception("You can only tweak functions.")

        return output

    @staticmethod
    def getsource(ob):
        """
            Get's the source code of a function. If you aren't sure that you
            have the most basic bit of the function, use getsource() to find it
            before calling this.
        """
        global tweak_history, tweak_originals
        name = None

        fun = Tweaker.getfunction(ob)
        if not fun in tweak_history:

            #        try:
                sourcefile = inspect.getsourcefile(fun)
                lines, start = inspect.findsource(fun)

                blockfinder = inspect.BlockFinder()
                try:
                    tokenize.tokenize(iter(lines[start:]).next, blockfinder.tokeneater)
                except (inspect.EndOfBlock, IndentationError):
                    pass
                end = blockfinder.last

                tweak_originals[fun] = (sourcefile, start, start+end)
                tweak_history[fun] = ["".join(lines[start:start+end])]
            #except:
            #    raise Exception("Could not find source, probably an interpreter-defined function?")

        return tweak_history[fun][-1]

    @staticmethod
    def setsource(ob, src):
        """
            Redefine a function using the new definition given.
            src must be the complete function body including the
            def funcname(....):
                a = b
                ...

            Will fail unless getsource(fun) has been done before
            This is to prevent edit-conflicts on apply.
        """
        global tweak_history
        fun = Tweaker.getfunction(ob)
        try:
            tweak_history[fun].append(src)
        except:
            raise Exception("Cannot setsource() without getsource()")

        #Find the function definition, probably ok to get fun.__name__, but possible for user to change name in tweak
        #Can't use .match() as there may be decorators
        m = re.search(r'\s*def\s+(\w+)\s*\(',src)
        if m:
            funname = m.group(1)
            interpreter = code.InteractiveConsole(fun.func_globals)

            #Strip whitespace from decorators and def
            lines = src.split("\n")

            for i in range(len(lines)):
                lines[i] = lines[i].lstrip()
                if lines[i].startswith("def"):
                    break

            src = "\n".join(lines)

            if not interpreter.push(src):
                newobj = interpreter.locals[funname]
                if inspect.ismethoddescriptor(newobj):
                    newobj = newobj.__get__(1) #I have no idea wht significance this paramter has.
                for attr in dir(newobj):
                    if attr not in ['func_closure', 'func_globals', '__class__']:
                        setattr(fun, attr, getattr(newobj, attr))
                return
            else:
                raise Exception("True?")

        raise Exception("Could not setsource() probably invalid syntax.")

    @staticmethod
    def apply(ob,filename=None):
        """
            Saves the changes you have made during this lot of tweaking to the file.
            TODO allow putting modules/classes with modified functions here.

            By default the filename will be that which was read from, but in some cases
            it may be useful to override that.
        """
        global tweak_originals
        global tweak_history
        fun = Tweaker.getfunction(ob)

        try:
            (sourcefile, start, end) = tweak_originals[fun]
            if filename:
                sourcefile = filename
        except:
            return #Not tweaked, so already applied

        lines = linecache.getlines(sourcefile) #The file has already been read into python, so no need to read it again

        if not "".join(lines[start:end]) == tweak_history[fun][0]:
            raise Exception("Source file has changed on disk")

        shutil.copy(sourcefile, sourcefile + ".bak" + str(int(time.time())) )
        file = open(sourcefile, "w")
        file.write("".join(lines[:start]) + tweak_history[fun][-1] + "".join(lines[end:]))
        file.close()
        print "Written new version"

    @staticmethod
    def revert(ob, steps=1):
        global tweak_history
        fun = Tweaker.getfunction(ob)
        if fun in tweak_history:
            setsource(fun, tweak_history[fun][0-steps])
        else:
            raise Exception("%s has not been tweaked")


T = Tweaker()
tweak = T.tweak
apply = T.apply
revert = T.revert
getsource = T.getsource
setsource = T.setsource
getfunction = T.getfunction