Update contrib/python/ipython/py3 to 8.15.0

3 files changed, 421 insertions, 9 deletions

diff --git a/contrib/python/ipython/py3/IPython/core/magics/ast_mod.py b/contrib/python/ipython/py3/IPython/core/magics/ast_mod.py
new file mode 100644
index 0000000000..e28b9f1231
--- /dev/null
+++ b/contrib/python/ipython/py3/IPython/core/magics/ast_mod.py
@@ -0,0 +1,320 @@
+"""
+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)
diff --git a/contrib/python/ipython/py3/IPython/core/magics/basic.py b/contrib/python/ipython/py3/IPython/core/magics/basic.py
index 814dec72e2..54b0c2a4de 100644
--- a/contrib/python/ipython/py3/IPython/core/magics/basic.py
+++ b/contrib/python/ipython/py3/IPython/core/magics/basic.py
@@ -109,12 +109,12 @@ class BasicMagics(Magics):
           Created `%%t` as an alias for `%%timeit`.
 
           In [2]: %t -n1 pass
-          1 loops, best of 3: 954 ns per loop
+          107 ns ± 43.6 ns per loop (mean ± std. dev. of 7 runs, 1 loop each)
 
           In [3]: %%t -n1
              ...: pass
              ...:
-          1 loops, best of 3: 954 ns per loop
+          107 ns ± 58.3 ns per loop (mean ± std. dev. of 7 runs, 1 loop each)
 
           In [4]: %alias_magic --cell whereami pwd
           UsageError: Cell magic function `%%pwd` not found.
diff --git a/contrib/python/ipython/py3/IPython/core/magics/execution.py b/contrib/python/ipython/py3/IPython/core/magics/execution.py
index 228cbd9da7..4147dac963 100644
--- a/contrib/python/ipython/py3/IPython/core/magics/execution.py
+++ b/contrib/python/ipython/py3/IPython/core/magics/execution.py
@@ -8,6 +8,7 @@
 import ast
 import bdb
 import builtins as builtin_mod
+import copy
 import cProfile as profile
 import gc
 import itertools
@@ -19,14 +20,28 @@ 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 @@ 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 @@ 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 @@ 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)
@@ -1035,7 +1050,7 @@ class ExecutionMagics(Magics):
         provided, <N> is determined so as to get sufficient accuracy.
 
         -r<R>: number of repeats <R>, each consisting of <N> loops, and take the
-        best result.
+        average result.
         Default: 7
 
         -t: use time.time to measure the time, which is the default on Unix.
@@ -1474,6 +1489,83 @@ 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 @@ 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])