Injecting code into a function

[Bengt Richter]

On 25 Apr 2005 23:31:29 -0700, "George Sakkis" <gsakkis at rutgers.edu> wrote:

>> 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).
>
>That's all good, at least if I knew how to poke with bytecodes ;-)
>What's a good starting point to look at ?
>
>By the way, the original problem was yet another property packager,
>which I posted as recipe at the cookbook:
>http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/410698.
>
I made a couple of byte code hacking decorators based on the techniques
in Raymond Hettinger's optimization decorators, which optimize access to
globals and fold constant expressions and I forgot what else.

I stripped all the optimization out since I was interested in injecting
code to preset values for selected named local variables, to avoid using
the default argument value hack. Another variant used this plus modifying
the call signature to do a version of currying. Also modified the line
number map. You will have some trickier stuff to work out though, I think,
since there may be many returns, and they can happen in try/except/finally
context, and you'll have to see how that translates (you can see with
disassembler), and then if you can't just overwrite some returns to jump
to a common place, you may have to widen a place to put code and adjust
any jumps across modified stuff.

It would be much nicer if you can get the source and do it via AST mods.
I am toying with the concept of a custom importer that does import-time
"decoration" of statements or statement sequences (with their suites)

I am thinking of using '@@' something like

    # ... statments
    @@module.decorator(nstatements, args) # how many statements, and whatever other args
    # decorated statement(s)
    ...   

Where module has to be imported already at the time of the decorating-importer's importation
of the above (simplest way to guarantee that it's available to execute, but not the only way).

Of course this means that a .py file has to be available for the code you are interesting in
modifying, and you have to catch the import. (Since there's a '@@' you will know if you don't ;-)

Anyway, I am not going to have time for that for some time, so I'll mention what I had in mind,
in case someone wants to do something with it:

The decorating-importer would find the @@ lines and e.g. change

    @@decomodule.decorator(nstatements, <anything legal>) # how many statements, and whatever other args

to

    __AT_AT_DECO__ = decomodule.decorator(nstatementsliteral, <anything legal> ...)

which should just be a matter of src.replace('@@','__AT_AT_DECO__ = ') on lines where
line.lstrip().startswith('@@') is true, thus making valid python source.

It would then call compiler.parse with the modified source for an AST to work with.
It would walk the AST to find the __AT_AT_DECO__ assignments (saving a reference to the containing statement
list and index of the __AT_AT_DECO__ = ... statement in the list, and extract the nstatementsliteral
and then call the decomodule.decorator with a slice of statements to be replaced, including the __AT_AT_DECO__
statement, which is passed as the first statement AST node in the slice, and may have interesting things
in the <anything legal> part, which the decorating-importer doesn't need to know about. It just needs the
two names decomodule and decorator and the number of statements, replacing them something like

    containing_statement_list[index:index+int(nstatementsliteral)+1] = getattr(__import__(decomodule), decorator)(
             containing_statement_list[index:index+int(nstatementsliteral)+1])

IOW pass the slice of statement list nodes to the decorator and expect it to get additional info from the
__AT_AT_DECO__ = statement that is included, and expect it to return a replacement list of AST statment nodes to
replace what it got, which might be minor or major modifications.

The walk would continue until all the __AT_AT_DECO__ regions were replaced, and then the whole AST would be
compiled and made into a module delivered as if it were an ordinary import.

Note that this is not a macro preprocessor, but depending on what is in decomodule.decorator, one could mess with
legal source (after src.replace('@@','__AT_AT_DECO__ = ') it has to satisfy compiler.parse, and later check also)
pretty seriously, but it has to be pretty legal to start. Of course one could abuse magic-content strings, but
metaclasses can already do a lot of that. That's not what I'm hoping to do. Just super-duper-decoration ;-)

Anyway, that's a sketch of the idea. I won't have time for it or much else for a longish while now.
If someone's interested, they can quote this and start another thread ;-)

Regards,
Bengt Richter