IPython.core.magics.ast_mod

This module contains utility function and classes to inject simple ast transformations based on code strings into IPython. While it is already possible with ast-transformers it is not easy to directly manipulate ast.

IPython has pre-code and post-code hooks, but are ran from within the IPython machinery so may be inappropriate, for example for performance measurement.

This module give you tools to simplify this, and expose 2 classes:

• ReplaceCodeTransformer which is a simple ast transformer based on code template,

and for advance case:

• Mangler which is a simple ast transformer that mangle names in the ast.

Example, let's try to make a simple version of the timeit magic, that run a code snippet 10 times and print the average time taken.

Basically we want to run :

from time import perf_counter
now = perf_counter()
for i in range(10):
    __code__ # our code
print(f"Time taken: {(perf_counter() - now)/10}")
__ret__ # the result of the last statement

Where __code__ is the code snippet we want to run, and __ret__ is the result, so that if we for example run dataframe.head() IPython still display the head of dataframe instead of nothing.

Here is a complete example of a file timit2.py that define such a magic:

from IPython.core.magic import (
    Magics,
    magics_class,
    line_cell_magic,
)
from IPython.core.magics.ast_mod import ReplaceCodeTransformer
from textwrap import dedent
import ast

template = template = dedent('''
    from time import perf_counter
    now = perf_counter()
    for i in range(10):
        __code__
    print(f"Time taken: {(perf_counter() - now)/10}")
    __ret__
'''
)


@magics_class
class AstM(Magics):
    @line_cell_magic
    def t2(self, line, cell):
        transformer = ReplaceCodeTransformer.from_string(template)
        transformer.debug = True
        transformer.mangler.debug = True
        new_code = transformer.visit(ast.parse(cell))
        return exec(compile(new_code, "<ast>", "exec"))


def load_ipython_extension(ip):
    ip.register_magics(AstM)

In [1]: %load_ext timit2

In [2]: %%t2
   ...: import time
   ...: time.sleep(0.05)
   ...:
   ...:
Time taken: 0.05435649999999441

If you wish to ran all the code enter in IPython in an ast transformer, you can do so as well:

In [1]: from IPython.core.magics.ast_mod import ReplaceCodeTransformer
   ...:
   ...: template = '''
   ...: from time import perf_counter
   ...: now = perf_counter()
   ...: __code__
   ...: print(f"Code ran in {perf_counter()-now}")
   ...: __ret__'''
   ...:
   ...: get_ipython().ast_transformers.append(ReplaceCodeTransformer.from_string(template))

In [2]: 1+1
Code ran in 3.40410006174352e-05
Out[2]: 2

Hygiene and Mangling

The ast transformer above is not hygienic, it may not work if the user code use the same variable names as the ones used in the template. For example.

To help with this by default the ReplaceCodeTransformer will mangle all names staring with 3 underscores. This is a simple heuristic that should work in most case, but can be cumbersome in some case. We provide a Mangler class that can be overridden to change the mangling heuristic, or simply use the mangle_all utility function. It will _try_ to mangle all names (except __ret__ and __code__), but this include builtins (print, range, type) and replace those by invalid identifiers py prepending mangle-: mangle-print, mangle-range, mangle-type etc. This is not a problem as currently Python AST support invalid identifiers, but it may not be the case in the future.

You can set ReplaceCodeTransformer.debug=True and ReplaceCodeTransformer.mangler.debug=True to see the code after mangling and transforming:

In [1]: from IPython.core.magics.ast_mod import ReplaceCodeTransformer, mangle_all
   ...:
   ...: template = '''
   ...: from builtins import type, print
   ...: from time import perf_counter
   ...: now = perf_counter()
   ...: __code__
   ...: print(f"Code ran in {perf_counter()-now}")
   ...: __ret__'''
   ...:
   ...: transformer = ReplaceCodeTransformer.from_string(template, mangling_predicate=mangle_all)


In [2]: transformer.debug = True
   ...: transformer.mangler.debug = True
   ...: get_ipython().ast_transformers.append(transformer)

In [3]: 1+1
Mangling Alias mangle-type
Mangling Alias mangle-print
Mangling Alias mangle-perf_counter
Mangling now
Mangling perf_counter
Not mangling __code__
Mangling print
Mangling perf_counter
Mangling now
Not mangling __ret__
---- Transformed code ----
from builtins import type as mangle-type, print as mangle-print
from time import perf_counter as mangle-perf_counter
mangle-now = mangle-perf_counter()
ret-tmp = 1 + 1
mangle-print(f'Code ran in {mangle-perf_counter() - mangle-now}')
ret-tmp
---- ---------------- ----
Code ran in 0.00013654199938173406
Out[3]: 2