Provide and easier way to generate magics and pre-post hooks

This introduce a code base way of modifying the ast.
This allow to use Template strings with the two special variable names
names `__code__` and `__ret__` allowing to quickly write a magic, or
hooks that modify the ast.

This also introduce a `%code_wrap` cell magic to on the fly wrap code.
It is this easy to for example modify IPython to say time each block of
code, or retry them, or wrap them in try/except and analyse the error
message, profile...

Note that this is not new, but simply convenience function and
utilities, especially around hygiene.

IPython/core/magics/ast_mod.py

@@ -0,0 +1,320 b''
+"""
+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)

IPython/core/interactiveshell.py

@@ -3338,8 +3338,11 b' class InteractiveShell(SingletonConfigurable):'
                 # 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/execution.py

@@ -8,6 +8,7 b''
 import ast
 import bdb
 import builtins as builtin_mod
+import copy
 import cProfile as profile
 import gc
 import itertools
@@ -19,14 +20,28 b' import shlex'
 import sys
 import time
 import timeit
-from ast import Module
+from typing import Dict, Any
+from ast import (
+    Assign,
+    Call,
+    Expr,
+    Load,
+    Module,
+    Name,
+    NodeTransformer,
+    Store,
+    parse,
+    unparse,
+)
 from io import StringIO
 from logging import error
 from pathlib import Path
 from pdb import Restart
+from textwrap import dedent, indent
 from warnings import warn
 
 from IPython.core import magic_arguments, oinspect, page
+from IPython.core.displayhook import DisplayHook
 from IPython.core.error import UsageError
 from IPython.core.macro import Macro
 from IPython.core.magic import (
@@ -37,8 +52,8 b' from IPython.core.magic import ('
     magics_class,
     needs_local_scope,
     no_var_expand,
-    output_can_be_silenced,
     on_off,
+    output_can_be_silenced,
 )
 from IPython.testing.skipdoctest import skip_doctest
 from IPython.utils.capture import capture_output
@@ -47,7 +62,7 b' from IPython.utils.ipstruct import Struct'
 from IPython.utils.module_paths import find_mod
 from IPython.utils.path import get_py_filename, shellglob
 from IPython.utils.timing import clock, clock2
-from IPython.core.displayhook import DisplayHook
+from IPython.core.magics.ast_mod import ReplaceCodeTransformer
 
 #-----------------------------------------------------------------------------
 # Magic implementation classes
@@ -164,9 +179,9 b' class Timer(timeit.Timer):'
 
 @magics_class
 class ExecutionMagics(Magics):
-    """Magics related to code execution, debugging, profiling, etc.
+    """Magics related to code execution, debugging, profiling, etc."""
 
-    """
+    _transformers: Dict[str, Any] = {}
 
     def __init__(self, shell):
         super(ExecutionMagics, self).__init__(shell)
@@ -1474,6 +1489,83 b' class ExecutionMagics(Magics):'
         elif args.output:
             self.shell.user_ns[args.output] = io
 
+    @skip_doctest
+    @magic_arguments.magic_arguments()
+    @magic_arguments.argument("name", type=str, default="default", nargs="?")
+    @magic_arguments.argument(
+        "--remove", action="store_true", help="remove the current transformer"
+    )
+    @magic_arguments.argument(
+        "--list", action="store_true", help="list existing transformers name"
+    )
+    @magic_arguments.argument(
+        "--list-all",
+        action="store_true",
+        help="list existing transformers name and code template",
+    )
+    @line_cell_magic
+    def code_wrap(self, line, cell=None):
+        """
+        Simple magic to quickly define a code transformer for all IPython's future imput.
+
+        ``__code__`` and ``__ret__`` are special variable that represent the code to run
+        and the value of the last expression of ``__code__`` respectively.
+
+        Examples
+        --------
+
+        .. ipython::
+
+            In [1]: %%code_wrap before_after
+               ...: print('before')
+               ...: __code__
+               ...: print('after')
+               ...: __ret__
+
+
+            In [2]: 1
+            before
+            after
+            Out[2]: 1
+
+            In [3]: %code_wrap --list
+            before_after
+
+            In [4]: %code_wrap --list-all
+            before_after :
+                print('before')
+                __code__
+                print('after')
+                __ret__
+
+            In [5]: %code_wrap --remove before_after
+
+        """
+        args = magic_arguments.parse_argstring(self.code_wrap, line)
+
+        if args.list:
+            for name in self._transformers.keys():
+                print(name)
+            return
+        if args.list_all:
+            for name, _t in self._transformers.items():
+                print(name, ":")
+                print(indent(ast.unparse(_t.template), "    "))
+            print()
+            return
+
+        to_remove = self._transformers.pop(args.name, None)
+        if to_remove in self.shell.ast_transformers:
+            self.shell.ast_transformers.remove(to_remove)
+        if cell is None or args.remove:
+            return
+
+        _trs = ReplaceCodeTransformer(ast.parse(cell))
+
+        self._transformers[args.name] = _trs
+        self.shell.ast_transformers.append(_trs)
+
+
 def parse_breakpoint(text, current_file):
     '''Returns (file, line) for file:line and (current_file, line) for line'''
     colon = text.find(':')
@@ -1519,4 +1611,4 b' def _format_time(timespan, precision=3):'
         order = min(-int(math.floor(math.log10(timespan)) // 3), 3)
     else:
         order = 3
-    return u"%.*g %s" % (precision, timespan * scaling[order], units[order])
+    return "%.*g %s" % (precision, timespan * scaling[order], units[order])

docs/autogen_api.py

@@ -24,8 +24,6 b" if __name__ == '__main__':"
     docwriter.package_skip_patterns += [r'\.external$',
                                         # Extensions are documented elsewhere.
                                         r'\.extensions',
-                                        # Magics are documented separately
-                                        r'\.core\.magics',
                                         # This isn't API
                                         r'\.sphinxext',
                                         # Shims