Modify function in decorator

Question

I was thinking about making a decorator for the purpose of increasing performance. A decorator that modifies the source code of the function it decorates, and returns the modified function.

While thinking this through, I figured that if I could just get the source code of the function, I could do this. But is it possible to access the source code of a function inside a decorator? If I have a decorator like this:

import inspect

def decorate(f):
    exec(inspect.getsource(f))
    return eval(f.__name__)

@decorate
def test():
    return 1

I get an OSError:

OSError: could not get source code

This appears to be because test is not fully formed before it is passed into decorate. However, this works:

import inspect

def decorate(f):
    exec(inspect.getsource(f))
    return eval(f.__name__)

def test():
    return 1
test = decorate(test)

It just doesn't have that decorator flair to it, though. It seems that this might be possible, because f.__code__ is defined.

---

Upon further inspection, it appears that this only happens when I put the inspect.getsource(f) into exec. Otherwise, it seems that I can get the source code.

---

As a rough sketch of the first thing that's on my mind, I'm thinking of tail-recursion. I wrote this decorator that is unfortunately slow and requires a very specific style of writing the function to be decorated:

def tail_recurse(acc_default):
    def decorate(f):
        def wrapper(*args, acc=acc_default):
            args = args + (acc,)
            while True:
                return_type, *rargs = f(*args)
                if return_type is None:
                    return rargs[-1]
                args = rargs
        return wrapper
    return decorate

Basically, I'm thinking of doing something as simple as replacing the body of a function with:

while True:
    __body__
    update_args

Answer 1

You can use functools.wraps with your original code:

import inspect
from functools import wraps

@wraps
def decorate(f):
    exec(inspect.getsource(f))
    return eval(f.__name__)

@decorate
def test():
    return 1

Output:

In [2]: test()
Out[2]: 1

If you plan on changing source at runtime then you should get familiar with the ast library, there is an excellent video from pycon 2011 where Matthew Desmarais gives a talk on how to use the ast module to change the source code from the basics right up to more the more advanced options, this is a simple working example of the python to javascript translator that is used in the talk, it will work for simple examples like the fib function provided.

It should give you a good understanding of how the NodeTransformer works which is what you will want to use to manipulate your code at runtime, you can decorate your functions using something similar to the dec function below, the difference will be you will be returning compiled code:

from ast import parse, NodeTransformer


class Transformer(NodeTransformer):
    def __init__(self):
        self.src = ""
        self.indent = 0

    def translate(self, node):
        self.visit(node)
        return self.src

    def _indent(self, line):
        return "{}{line}".format(" " * self.indent, line=line)

    def render(self, body):
        self.indent += 2
        for stmt in body:
            self.visit(stmt)
        self.indent -= 2

    def visit_Num(self, node):
        self.src += "{}".format(node.n)

    def visit_Str(self, node):
        self.src += "{}".format(node.s)

    def visit_FunctionDef(self, defn):
        args = ",".join(name.arg for name in defn.args.args)
        js_defn = "var {} = function({}){{\n"
        self.src += self._indent(js_defn.format(defn.name, args))
        self.render(defn.body)
        self.src += self._indent("}\n")

    def visit_Eq(self, less):
        self.src += "=="

    def visit_Name(self, name):
        self.src += "{}".format(name.id)

    def visit_BinOp(self, binop):
        self.visit(binop.left)
        self.src += " "
        self.visit(binop.op)
        self.src += " "
        self.visit(binop.right)

    def visit_If(self, _if):
        self.src += self._indent("if (")
        self.visit(_if.test)
        self.src += ") {\n"
        self.render(_if.body)
           self.src += " "*self.indent + "}\n"


    def visit_Compare(self, comp):
        self.visit(comp.left)
        self.src += " "
        self.visit(comp.ops[0])
        self.src += " "
        self.visit(comp.comparators[0])

    def visit_Call(self, call):
        self.src += " "
        self.src += "{}(".format(call.func.id)
        self.visit(call.args[0])
        self.src += ")"

    def visit_Add(self, add):
        self.src += "+"

    def visit_Sub(self, add):
        self.src += "-"

    def visit_Return(self, ret):
        self.src += self._indent("return")
        if ret.value:
            self.src += " "
            self.visit(ret.value)
        self.src += ";\n"


def dec(f):
    source = getsource(f)
    _ast = parse(source)
    trans = Transformer()
    trans.indent = 0
    return trans.translate(_ast)


from inspect import getsource


def fibonacci(n):
    if n == 0:
        return 0
    if n == 1:
        return 1
    return fibonacci(n - 1) + fibonacci(n - 2)

Running the dec function outputs our python as javascript:

print(dec(fibonacci))
var fibonacci = function(n){
  if (n == 0) {
    return 0;
  }
  if (n == 1) {
    return 1;
  }
  return  fibonacci(n - 1) +  fibonacci(n - 2);
}

The greentreesnakes docs are also worth a read.

Answer 2

This works:

import inspect, itertools

def decorate(f):
    source = itertools.dropwhile(lambda line: line.startswith('@'), inspect.getsource(f).splitlines())
    exec('\n'.join(source))
    return eval(f.__name__)

@decorate
def test():
    return 1

I think the problem is the inclusion of the decorator in the function source.

# foo.py
import inspect

def decorate(f):
    print inspect.getsource(f)

@decorate
def test():
    return 1

 

>>> import foo
@decorate
def test():
    return 1
>>> # Notice the decorator is included in the source.

exec sees @decorate for a test defined in a string, so it calls decorate recursively, but this time inspect.getsource fails, because it can't find the source of a function defined in a string.