transform python AST

ast_transformer.py

"""
In Public domain see https://gist.github.com/Carreau/0f051f57734222da925cd364e59cc17e
"""
from ast import NodeTransformer, Expr ,parse, Assign, Name, Store, Call, Load
import ast
from textwrap import dedent

class Mangler(NodeTransformer):
    """
    Mangle given names in and ast tree to make sure they do not conflict with user code.    
    """

    def __init__(self, mangle):
        self.mangle = mangle
        
    def visit_Name(self, node):
        if node.id in self.mangle:
            node.id = 'mangle-'+node.id
        return node
        

class PrePostAstTransformer(NodeTransformer):
    """
    Allow to safely wrap user code with pre/post execution hooks that
    run just before and just after usercode, __inside__ the execution loop, 
    But still returns the value of the last Expression. 
    
    This might not behave as expected if the user change the InteractiveShell.ast_node_interactivity option.
    
    This is currently not hygienic and care must be taken to use uncommon names in the pre/post block.
    
    Assuming the user have 
    
    ```
    code_block:
        [with many expressions]
    last_expression
    ```
    
    It will transform it into
    
    ```
    try:
        pre_block
        code_block:
            [with many expressions]
        return_value = last_expression
    finally:
        post_block
    return_value
    ```
    
    Thus makind sure that post is always executed even if pre or user code fails.    
    """

    def __init__(self, pre, post):
        """
        pre and post are either strings, or ast.Modules object that need to be run just before or after
        the user code.
        
        While strings are possible, we suggest using ast.Modules
        object and mangling the corresponding varaibles names
        to be invalid python identifiers to avoid name conflicts.
        """
        if isinstance(pre, str):
            pre = parse(pre)
        if isinstance(post, str):
            pre = parse(post)
            
        self.pre = pre.body
        self.post = post.body
        self.active = True
        
    def reset(self):

        self.core = parse(dedent("""
            try:
                pass
            finally:
                pass
            """))
        self.try_ = self.core.body[0].body = []
        self.fin = self.core.body[0].finalbody = []
        
    def visit_Module(self, node):
        if not self.active:
            return node
        self.reset()
        last = node.body[-1]
        ret = None
        if isinstance(last, Expr):
            node.body.pop()
            node.body.append(Assign([Name('ast-tmp', ctx=Store())], value=last.value ))
            ret = Expr(value=Name('ast-tmp', ctx=Load()))
        #self.core.body.insert(0, Assign([Name('_p', ctx=Store())], value=ast.Constant(None) ))
        if ret:
            self.core.body.insert(0, Assign([Name('ast-tmp', ctx=Store())], value=ast.Constant(None) ))
        for p in self.pre+node.body:
            self.try_.append(p)
        for p in self.post:
            self.fin.append(p)
        if ret is not None:
            self.core.body.append(ret)

        ast.fix_missing_locations(self.core)
        return self.core
        
mangle = Mangler(['_p', '_PX']).visit        
pre = mangle(parse("""
from pyinstrument import Profiler
from IPython.display import HTML
_p = Profiler()
_p.start()
"""))

post = mangle(parse("""
if _p is not None:
    _p.stop()
    text = _p.output_html()
    call_id = 1
    # This should just be fixed at the source or via an iframe...
    app_id = f"app-{call_id}"
    text = text.replace('id="app"', f'id="{app_id}"')
    text = text.replace("#app", f"#{app_id}")
    text = text.replace("document.title", "{}")
    text = text.replace("this.setFavicon", "this.setFavicon || function(){}")
    display(HTML(text))
"""))


ip = get_ipython()
ip.ast_transformers = [PrePostAstTransformer(pre, post)]

LICENSE

This is free and unencumbered software released into the public domain.

Anyone is free to copy, modify, publish, use, compile, sell, or
distribute this software, either in source code form or as a compiled
binary, for any purpose, commercial or non-commercial, and by any
means.

In jurisdictions that recognize copyright laws, the author or authors
of this software dedicate any and all copyright interest in the
software to the public domain. We make this dedication for the benefit
of the public at large and to the detriment of our heirs and
successors. We intend this dedication to be an overt act of
relinquishment in perpetuity of all present and future rights to this
software under copyright law.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
OTHER DEALINGS IN THE SOFTWARE.

For more information, please refer to <https://unlicense.org>