Python: monkey patch a function's source code

Question

Can I add a prefix and suffix to the source code of functions?

I know about decorators and do not want to use them (the minimal example below doesn't make clear why, but I have my reasons).

def f():
    print('world')
g = patched(f,prefix='print("Hello, ");',suffix='print("!");')
g() # Hello, world!

Here is what I have so far:

import inspect
import ast
import copy
def patched(f,prefix,suffix):
    source = inspect.getsource(f)
    tree = ast.parse(source)
    new_body = [
        ast.parse(prefix).body[0],
        *tree.body[0].body,
        ast.parse(suffix).body[0]
    ]
    tree.body[0].body = new_body
    g = copy.deepcopy(f)
    g.__code__ = compile(tree,g.__code__.co_filename,'exec')
    return g

Unfortunately, nothing happens if I use this and then call g() as above; neither world nor Hello, world! are printed.

Answer 1

Here is a rough version of what can be done:

import inspect
import ast
import copy
def patched(f,prefix,suffix):
    source = inspect.getsource(f)
    tree = ast.parse(source)
    new_body = [
        ast.parse(prefix).body[0],
        *tree.body[0].body,
        ast.parse(suffix).body[0]
    ]
    tree.body[0].body = new_body
    code = compile(tree,filename=f.__code__.co_filename,mode='exec')
    namespace = {}
    exec(code,namespace)
    g = namespace[f.__name__]
    return g

def temp():
    pass
def f():
    print('world',end='')
g = patched(f,prefix='print("Hello, ",end="")',suffix='print("!",end="")')
g() # Hello, world!

The call of compile compiles an entire module (represented by tree). This module is then executed in an empty namespace from which the desired function is finally extracted. (Warning: the namespace will need to be filled with some globals from where f comes from if f uses those.)

---

After some more work, here is a real example of what can be done with this. It uses some extended version of the principle above:

import numpy as np
from playground import graphexecute
@graphexecute(verbose=True)
def my_algorithm(x,y,z):
    def SumFirstArguments(x,y)->sumxy:
        sumxy = x+y
    def SinOfThird(z)->sinz:
        sinz = np.sin(z)
    def FinalProduct(sumxy,sinz)->prod:
        prod = sumxy*sinz
    def Return(prod):
        return prod
print(my_algorithm(x=1,y=2,z=3)) 
#OUTPUT:
#>>Executing part SumFirstArguments
#>>Executing part SinOfThird
#>>Executing part FinalProduct
#>>Executing part Return
#>>0.4233600241796016

The clou is that I get the exact same output if I reshuffle the parts of my_algorithm, for example like this:

@graphexecute(verbose=True)
def my_algorithm2(x,y,z):
    def FinalProduct(sumxy,sinz)->prod:
        prod = sumxy*sinz
    def SumFirstArguments(x,y)->sumxy:
        sumxy = x+y
    def SinOfThird(z)->sinz:
        sinz = np.sin(z)
    def Return(prod):
        return prod
print(my_algorithm2(x=1,y=2,z=3)) 
#OUTPUT:
#>>Executing part SumFirstArguments
#>>Executing part SinOfThird
#>>Executing part FinalProduct
#>>Executing part Return
#>>0.4233600241796016

This works by (1) grabbing the source of my_algorithm and turning it into an ast (2) patching each function defined within my_algorithm (e.g. SumFirstArguments) to return locals (3) deciding based on the inputs and the outputs (as defined by the type hints) in which order the parts of my_algorithm should be executed. Furthermore, a possibility that I do not have implemented yet is to execute independent parts in parallel (such as SumFirstArguments and SinOfThird). Let me know if you want the sourcecode of graphexecute, I haven't included it here because it contains a lot of stuff that is not relevant to this question.