diff --git a/IPython/core/interactiveshell.py b/IPython/core/interactiveshell.py index 7392de7..3dfaecb 100644 --- a/IPython/core/interactiveshell.py +++ b/IPython/core/interactiveshell.py @@ -3338,8 +3338,11 @@ 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: diff --git a/IPython/core/magics/ast_mod.py b/IPython/core/magics/ast_mod.py new file mode 100644 index 0000000..e28b9f1 --- /dev/null +++ b/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, "", "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/IPython/core/magics/execution.py b/IPython/core/magics/execution.py index 228cbd9..f9942dd 100644 --- a/IPython/core/magics/execution.py +++ b/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) @@ -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]) diff --git a/docs/autogen_api.py b/docs/autogen_api.py index 0011900..10e9cbc 100755 --- a/docs/autogen_api.py +++ b/docs/autogen_api.py @@ -24,8 +24,6 @@ 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