macro magic

IPython/core/interactiveshell.py

@@ -3338,8 +3338,11 @@ def transform_ast(self, node):
                 # an InputRejected.  Short-circuit in this case so that we
                 # don't unregister the transform.
                 raise
-            except Exception:
-                warn("AST transformer %r threw an error. It will be unregistered." % transformer)
+            except Exception as e:
+                warn(
+                    "AST transformer %r threw an error. It will be unregistered. %s"
+                    % (transformer, e)
+                )
                 self.ast_transformers.remove(transformer)
 
         if self.ast_transformers:

IPython/core/magics/ast_mod.py

@@ -0,0 +1,318 @@
+"""
+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 mesurement.
+
+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 :
+
+.. code-block:: python
+    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:
+
+.. code-block:: python
+    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)
+
+
+
+.. code-block:: python
+
+    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:
+
+.. code-block:: python
+
+    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:
+
+.. code-block:: python
+
+
+    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
+
+
+"""
+
+__skip_doctest__ = True
+
+
+from ast import NodeTransformer, Store, Load, Name, Expr, Assign, Module
+import ast
+import copy
+
+from typing import Dict, Optional
+
+
+mangle_all = lambda name: False if name in ("__ret__", "__code__") else True
+
+
+class Mangler(NodeTransformer):
+    """
+    Mangle given names in and ast tree to make sure they do not conflict with
+    user code.
+    """
+
+    enabled: bool = True
+    debug: bool = False
+
+    def log(self, *args, **kwargs):
+        if self.debug:
+            print(*args, **kwargs)
+
+    def __init__(self, predicate=None):
+        if predicate is None:
+            predicate = lambda name: name.startswith("___")
+        self.predicate = predicate
+
+    def visit_Name(self, node):
+        if self.predicate(node.id):
+            self.log("Mangling", node.id)
+            # Once in the ast we do not need
+            # names to be valid identifiers.
+            node.id = "mangle-" + node.id
+        else:
+            self.log("Not mangling", node.id)
+        return node
+
+    def visit_FunctionDef(self, node):
+        if self.predicate(node.name):
+            self.log("Mangling", node.name)
+            node.name = "mangle-" + node.name
+        else:
+            self.log("Not mangling", node.name)
+
+        for arg in node.args.args:
+            if self.predicate(arg.arg):
+                self.log("Mangling function arg", arg.arg)
+                arg.arg = "mangle-" + arg.arg
+            else:
+                self.log("Not mangling function arg", arg.arg)
+        return self.generic_visit(node)
+
+    def visit_ImportFrom(self, node):
+        return self._visit_Import_and_ImportFrom(node)
+
+    def visit_Import(self, node):
+        return self._visit_Import_and_ImportFrom(node)
+
+    def _visit_Import_and_ImportFrom(self, node):
+        for alias in node.names:
+            asname = alias.name if alias.asname is None else alias.asname
+            if self.predicate(asname):
+                new_name: str = "mangle-" + asname
+                self.log("Mangling Alias", new_name)
+                alias.asname = new_name
+            else:
+                self.log("Not mangling Alias", alias.asname)
+        return node
+
+
+class ReplaceCodeTransformer(NodeTransformer):
+    enabled: bool = True
+    debug: bool = False
+    mangler: Mangler
+
+    def __init__(
+        self, template: Module, mapping: Optional[Dict] = None, mangling_predicate=None
+    ):
+        assert isinstance(mapping, (dict, type(None)))
+        assert isinstance(mangling_predicate, (type(None), type(lambda: None)))
+        assert isinstance(template, ast.Module)
+        self.template = template
+        self.mangler = Mangler(predicate=mangling_predicate)
+        if mapping is None:
+            mapping = {}
+        self.mapping = mapping
+
+    @classmethod
+    def from_string(
+        cls, template: str, mapping: Optional[Dict] = None, mangling_predicate=None
+    ):
+        return cls(
+            ast.parse(template), mapping=mapping, mangling_predicate=mangling_predicate
+        )
+
+    def visit_Module(self, code):
+        if not self.enabled:
+            return code
+        # if not isinstance(code, ast.Module):
+        # recursively called...
+        #    return generic_visit(self, code)
+        last = code.body[-1]
+        if isinstance(last, Expr):
+            code.body.pop()
+            code.body.append(Assign([Name("ret-tmp", ctx=Store())], value=last.value))
+            ast.fix_missing_locations(code)
+            ret = Expr(value=Name("ret-tmp", ctx=Load()))
+            ret = ast.fix_missing_locations(ret)
+            self.mapping["__ret__"] = ret
+        else:
+            self.mapping["__ret__"] = ast.parse("None").body[0]
+        self.mapping["__code__"] = code.body
+        tpl = ast.fix_missing_locations(self.template)
+
+        tx = copy.deepcopy(tpl)
+        tx = self.mangler.visit(tx)
+        node = self.generic_visit(tx)
+        node_2 = ast.fix_missing_locations(node)
+        if self.debug:
+            print("---- Transformed code ----")
+            print(ast.unparse(node_2))
+            print("---- ---------------- ----")
+        return node_2
+
+    # this does not work as the name might be in a list and one might want to extend the list.
+    # def visit_Name(self, name):
+    #     if name.id in self.mapping and name.id == "__ret__":
+    #         print(name, "in mapping")
+    #         if isinstance(name.ctx, ast.Store):
+    #             return Name("tmp", ctx=Store())
+    #         else:
+    #             return copy.deepcopy(self.mapping[name.id])
+    #     return name
+
+    def visit_Expr(self, expr):
+        if isinstance(expr.value, Name) and expr.value.id in self.mapping:
+            if self.mapping[expr.value.id] is not None:
+                return copy.deepcopy(self.mapping[expr.value.id])
+        return self.generic_visit(expr)