Search code examples
pythonbytecodeintrospection

Change reference to function in run-time in Python


I need to change a call to a function inside another function during run-time.

Consider the following code:

def now():
    print "Hello World!"

class Sim:
    def __init__(self, arg, msg):
        self.msg = msg
        self.func = arg
        self.patch(self.func)

    def now(self):
        print self.msg

    def run(self):
        self.func()

    def patch(self, func):
        # Any references to the global now() in func
        # are replaced with the self.now() method.

def myfunc():
    now()

Then ...

>>> a = Sim(myfunc, "Hello Locals #1")
>>> b = Sim(myfunc, "Hello Locals #2")
>>> b.run()
Hello Locals #2
>>> a.run()
Hello Locals #1

One user has written code, myfunc(), that makes a call to a globally defined function now() and I cannot edit it. But I want it to call the method of the Sim instance instead. So I would need to "patch" the myfunc() function in run-time.

How could I go about this?

One possible solution is to edit the bytecode as done here: http://web.archive.org/web/20140306210310/http://www.jonathon-vogel.com/posts/patching_function_bytecode_with_python but I'm wondering if there is an easier way.


Solution

  • This is trickier than it might seem. To do it you need to:

    • Create a dict subclass with a __missing__ method to hold your new value of now. The __missing__ method then grabs any items not in the dictionary from the usual globals() dict.

    • Create a new function object from the existing myfunc() function, keeping its code object but using the new dictionary you created for globals.

    • Assign the new function back into globals using its function name.

    Here's how:

    def now():
        print "Hello World!"
    
    class NowGlobals(dict):
        def __init__(self, now, globals):
            self["now"] = now
            self.globals = globals
        def __missing__(self, key):
            return self.globals[key]
    
    class Sim(object):
        def __init__(self, func):
            func = self.patch(func)
            self.func = func
        def now(self):
            print "Hello locals!"
        def patch(self, func):
            funcname   = func.__name__
            nowglobals = NowGlobals(self.now, func.func_globals)
            func = type(func)(func.func_code, nowglobals)
            globals()[funcname] = func
            return func
    
    def myfunc():
        now()
    
    sim = Sim(myfunc)
    myfunc()
    

    There is really no need to have it in a class, but I've kept it that way since that's the way you wrote it originally.

    If myfunc is in another module, you'll need to rewrite Sim.patch() to patch back into that namespace, e.g. module.myfunc = sim.patch(module.myfunc)