Injecting code into a function

[Bengt Richter]

On 25 Apr 2005 03:32:38 -0700, "George Sakkis" <gsakkis at rutgers.edu> wrote:

>Is there a general way of injecting code into a function, typically
>before and/or after the existing code ? I know that for most purposes,
>an OO solution, such as the template pattern, is a cleaner way to get
>the same effect, but it's not always applicable (e.g. if you have no
>control over the design and you are given a function to start with). In
>particular, I want to get access to the function's locals() just before
>it exits, i.e. something like:
>
>def analyzeLocals(func):
>    func_locals = {}
>    def probeFunc():
>        # insert func's code here
>        sys._getframe(1).f_locals["func_locals"].update(locals())
>    probeFunc()
>    # func_locals now contains func's locals
>
>So, how can I add func's code in probeFunc so that the injected code
>(the update line here) is always called before the function exits ?
>That is, don't just inject it lexically in the end of the function if
>there are more than one exit points. I guess a solution will involve a
>good deal bytecode hacking, on which i know very little; if there's a
>link to a (relatively) simple HOWTO, it would be very useful.
>
I'm not clear on what your real goal is, but if you just want a snapshot
of what locals() is just before exiting func, that could be done with
a byte-code-hacking decorator with usage looking something like

    #func defined before this
    func_locals = {}
    @getlocals(func, func_locals)
    def probefunc(): pass

which would make a probefunc function that would be identical to func
except that before exiting, it would do func_locals.update(locals()).
(you might want func_locals to be a list and do func_locals.append(locals())
in case func is recursive and you are interested in the all the locals).

Alternatively, if this is a debugging thing, you might want to look into
sys.settrace -- as in this thing I cobbled together (not tested beyond what you see ;-):

----< tracelocals.py >----------------------------------------------------
class TraceLocals(object):
    from sys import settrace
    def __init__(self, *names, **kw):
        self.names = set(names)
        self.func_locals = kw.get('func_locals', [])    # [(name,locals()), ...] tuples
    def _gwatch(self, frame, event, arg):
        """
        Global scope watcher. When a new scope is entered, returns the local
        scope watching method _lwatch to do the work for that.
        """
        if event=='call':
            name = frame.f_code.co_name   # name of executing scope
            if name in self.names: return self._lwatch # name is one whose locals we want
                
    def _lwatch(self, frame, event, arg):
        if event == 'return':
            self.func_locals.append((frame.f_code.co_name,frame.f_locals))
        else:
            return self._lwatch # keep watching for return event
            
    def on(self):
        """Set the system trace hook to enable tracing. """
        self.settrace(self._gwatch)
    def off(self):
        """Reset the system trace hook to disable tracing. """
        self.settrace(None)
    
def main(modname, entry, *names):
    print 'main(', modname, entry, names,')'
    tr = TraceLocals(*names)
    mod = __import__(modname)
    try:
        tr.on()
        getattr(mod, entry)()
    finally:
        tr.off()
    return tr.func_locals

def test():
    tr = TraceLocals(*'t1 t2 t3'.split())
    def t1():
        x ='t1'
    def t2(y=123):
        y*=2
    def t3():
        t1()
        t2()
        t2('hello ')
    try:
        tr.on()
        t3()
    finally:
        tr.off()
    for name, loc in tr.func_locals: print '%5s: %s' %(name, loc)

if __name__=='__main__':
    import sys
    args = sys.argv[1:]
    if not args:
        raise SystemExit(
            'Usage: python tracelocals.py (-test | module entry name+)\n'
        )
    if args[0]=='-test': test()
    else:
        print args
        func_locals = main(args[0], args[1], *args[2:])
        for name, loc in func_locals: print '%5s: %s' %(name, loc)
--------------------------------------------------------------------------
Test result:

[22:37] C:\pywk\clp\sakkis\tracelocals>py24 tracelocals.py -test
   t1: {'x': 't1'}
   t2: {'y': 246}
   t2: {'y': 'hello hello '}
   t3: {'t2': <function t2 at 0x02EE8ED4>, 't1': <function t1 at 0x02EE8E9C>}

Note that t3 is seeing t1 and t2 in its locals -- I think because they're
visible as cell vars in test. If you put t1-t3 in a separate module, you don't see it:

----< tmod.py >---------------
def t1():
    print '-- t1'
    x ='t1'
def t2(y=123):
    print '-- t2'
    y*=2
def t3():
    print '-- t3'
    t1()
    t2()
    t2('hello ')
-----------------------------

[22:42] C:\pywk\clp\sakkis\tracelocals>py24 tracelocals.py tmod t3 t1 t2 t3
['tmod', 't3', 't1', 't2', 't3']
main( tmod t3 ('t1', 't2', 't3') )
-- t3
-- t1
-- t2
-- t2
   t1: {'x': 't1'}
   t2: {'y': 246}
   t2: {'y': 'hello hello '}
   t3: {}

[22:46] C:\pywk\clp\sakkis\tracelocals>py24 tracelocals.py tmod t3 t2
['tmod', 't3', 't2']
main( tmod t3 ('t2',) )
-- t3
-- t1
-- t2
-- t2
   t2: {'y': 246}
   t2: {'y': 'hello hello '}

Notice that the -- side effects from all being called in the last, but only t2 being captured.

Maybe this will give you something to expand to your needs.

Regards,
Bengt Richter