Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2This module contains utility function and classes to inject simple ast
3transformations based on code strings into IPython. While it is already possible
4with ast-transformers it is not easy to directly manipulate ast.
7IPython has pre-code and post-code hooks, but are ran from within the IPython
8machinery so may be inappropriate, for example for performance measurement.
10This module give you tools to simplify this, and expose 2 classes:
12- `ReplaceCodeTransformer` which is a simple ast transformer based on code
13 template,
15and for advance case:
17- `Mangler` which is a simple ast transformer that mangle names in the ast.
20Example, let's try to make a simple version of the ``timeit`` magic, that run a
21code snippet 10 times and print the average time taken.
23Basically we want to run :
25.. code-block:: python
27 from time import perf_counter
28 now = perf_counter()
29 for i in range(10):
30 __code__ # our code
31 print(f"Time taken: {(perf_counter() - now)/10}")
32 __ret__ # the result of the last statement
34Where ``__code__`` is the code snippet we want to run, and ``__ret__`` is the
35result, so that if we for example run `dataframe.head()` IPython still display
36the head of dataframe instead of nothing.
38Here is a complete example of a file `timit2.py` that define such a magic:
40.. code-block:: python
42 from IPython.core.magic import (
43 Magics,
44 magics_class,
45 line_cell_magic,
46 )
47 from IPython.core.magics.ast_mod import ReplaceCodeTransformer
48 from textwrap import dedent
49 import ast
51 template = template = dedent('''
52 from time import perf_counter
53 now = perf_counter()
54 for i in range(10):
55 __code__
56 print(f"Time taken: {(perf_counter() - now)/10}")
57 __ret__
58 '''
59 )
62 @magics_class
63 class AstM(Magics):
64 @line_cell_magic
65 def t2(self, line, cell):
66 transformer = ReplaceCodeTransformer.from_string(template)
67 transformer.debug = True
68 transformer.mangler.debug = True
69 new_code = transformer.visit(ast.parse(cell))
70 return exec(compile(new_code, "<ast>", "exec"))
73 def load_ipython_extension(ip):
74 ip.register_magics(AstM)
78.. code-block:: python
80 In [1]: %load_ext timit2
82 In [2]: %%t2
83 ...: import time
84 ...: time.sleep(0.05)
85 ...:
86 ...:
87 Time taken: 0.05435649999999441
90If you wish to ran all the code enter in IPython in an ast transformer, you can
91do so as well:
93.. code-block:: python
95 In [1]: from IPython.core.magics.ast_mod import ReplaceCodeTransformer
96 ...:
97 ...: template = '''
98 ...: from time import perf_counter
99 ...: now = perf_counter()
100 ...: __code__
101 ...: print(f"Code ran in {perf_counter()-now}")
102 ...: __ret__'''
103 ...:
104 ...: get_ipython().ast_transformers.append(ReplaceCodeTransformer.from_string(template))
106 In [2]: 1+1
107 Code ran in 3.40410006174352e-05
108 Out[2]: 2
112Hygiene and Mangling
113--------------------
115The ast transformer above is not hygienic, it may not work if the user code use
116the same variable names as the ones used in the template. For example.
118To help with this by default the `ReplaceCodeTransformer` will mangle all names
119staring with 3 underscores. This is a simple heuristic that should work in most
120case, but can be cumbersome in some case. We provide a `Mangler` class that can
121be overridden to change the mangling heuristic, or simply use the `mangle_all`
122utility function. It will _try_ to mangle all names (except `__ret__` and
123`__code__`), but this include builtins (``print``, ``range``, ``type``) and
124replace those by invalid identifiers py prepending ``mangle-``:
125``mangle-print``, ``mangle-range``, ``mangle-type`` etc. This is not a problem
126as currently Python AST support invalid identifiers, but it may not be the case
127in the future.
129You can set `ReplaceCodeTransformer.debug=True` and
130`ReplaceCodeTransformer.mangler.debug=True` to see the code after mangling and
131transforming:
133.. code-block:: python
136 In [1]: from IPython.core.magics.ast_mod import ReplaceCodeTransformer, mangle_all
137 ...:
138 ...: template = '''
139 ...: from builtins import type, print
140 ...: from time import perf_counter
141 ...: now = perf_counter()
142 ...: __code__
143 ...: print(f"Code ran in {perf_counter()-now}")
144 ...: __ret__'''
145 ...:
146 ...: transformer = ReplaceCodeTransformer.from_string(template, mangling_predicate=mangle_all)
149 In [2]: transformer.debug = True
150 ...: transformer.mangler.debug = True
151 ...: get_ipython().ast_transformers.append(transformer)
153 In [3]: 1+1
154 Mangling Alias mangle-type
155 Mangling Alias mangle-print
156 Mangling Alias mangle-perf_counter
157 Mangling now
158 Mangling perf_counter
159 Not mangling __code__
160 Mangling print
161 Mangling perf_counter
162 Mangling now
163 Not mangling __ret__
164 ---- Transformed code ----
165 from builtins import type as mangle-type, print as mangle-print
166 from time import perf_counter as mangle-perf_counter
167 mangle-now = mangle-perf_counter()
168 ret-tmp = 1 + 1
169 mangle-print(f'Code ran in {mangle-perf_counter() - mangle-now}')
170 ret-tmp
171 ---- ---------------- ----
172 Code ran in 0.00013654199938173406
173 Out[3]: 2
176"""
178__skip_doctest__ = True
181from ast import (
182 NodeTransformer,
183 Store,
184 Load,
185 Name,
186 Expr,
187 Assign,
188 Module,
189 Import,
190 ImportFrom,
191)
192import ast
193import copy
196mangle_all = lambda name: name not in ("__ret__", "__code__")
199class Mangler(NodeTransformer):
200 """
201 Mangle given names in and ast tree to make sure they do not conflict with
202 user code.
203 """
205 enabled: bool = True
206 debug: bool = False
208 def log(self, *args, **kwargs):
209 if self.debug:
210 print(*args, **kwargs)
212 def __init__(self, predicate=None):
213 if predicate is None:
214 predicate = lambda name: name.startswith("___")
215 self.predicate = predicate
217 def visit_Name(self, node):
218 if self.predicate(node.id):
219 self.log("Mangling", node.id)
220 # Once in the ast we do not need
221 # names to be valid identifiers.
222 node.id = "mangle-" + node.id
223 else:
224 self.log("Not mangling", node.id)
225 return node
227 def visit_FunctionDef(self, node):
228 if self.predicate(node.name):
229 self.log("Mangling", node.name)
230 node.name = "mangle-" + node.name
231 else:
232 self.log("Not mangling", node.name)
234 for arg in node.args.args:
235 if self.predicate(arg.arg):
236 self.log("Mangling function arg", arg.arg)
237 arg.arg = "mangle-" + arg.arg
238 else:
239 self.log("Not mangling function arg", arg.arg)
240 return self.generic_visit(node)
242 def visit_ImportFrom(self, node: ImportFrom):
243 return self._visit_Import_and_ImportFrom(node)
245 def visit_Import(self, node: Import):
246 return self._visit_Import_and_ImportFrom(node)
248 def _visit_Import_and_ImportFrom(self, node: Import | ImportFrom):
249 for alias in node.names:
250 asname = alias.name if alias.asname is None else alias.asname
251 if self.predicate(asname):
252 new_name: str = "mangle-" + asname
253 self.log("Mangling Alias", new_name)
254 alias.asname = new_name
255 else:
256 self.log("Not mangling Alias", alias.asname)
257 return node
260class ReplaceCodeTransformer(NodeTransformer):
261 enabled: bool = True
262 debug: bool = False
263 mangler: Mangler
265 def __init__(
266 self, template: Module, mapping: dict | None = None, mangling_predicate=None
267 ):
268 assert isinstance(mapping, (dict, type(None)))
269 assert isinstance(mangling_predicate, (type(None), type(lambda: None)))
270 assert isinstance(template, ast.Module)
271 self.template = template
272 self.mangler = Mangler(predicate=mangling_predicate)
273 if mapping is None:
274 mapping = {}
275 self.mapping = mapping
277 @classmethod
278 def from_string(
279 cls, template: str, mapping: dict | None = None, mangling_predicate=None
280 ):
281 return cls(
282 ast.parse(template), mapping=mapping, mangling_predicate=mangling_predicate
283 )
285 def visit_Module(self, code):
286 if not self.enabled:
287 return code
288 # if not isinstance(code, ast.Module):
289 # recursively called...
290 # return generic_visit(self, code)
291 last = code.body[-1]
292 if isinstance(last, Expr):
293 code.body.pop()
294 code.body.append(Assign([Name("ret-tmp", ctx=Store())], value=last.value))
295 ast.fix_missing_locations(code)
296 ret = Expr(value=Name("ret-tmp", ctx=Load()))
297 ret = ast.fix_missing_locations(ret)
298 self.mapping["__ret__"] = ret
299 else:
300 self.mapping["__ret__"] = ast.parse("None").body[0]
301 self.mapping["__code__"] = code.body
302 tpl = ast.fix_missing_locations(self.template)
304 tx = copy.deepcopy(tpl)
305 tx = self.mangler.visit(tx)
306 node = self.generic_visit(tx)
307 node_2 = ast.fix_missing_locations(node)
308 if self.debug:
309 print("---- Transformed code ----")
310 print(ast.unparse(node_2))
311 print("---- ---------------- ----")
312 return node_2
314 # this does not work as the name might be in a list and one might want to extend the list.
315 # def visit_Name(self, name):
316 # if name.id in self.mapping and name.id == "__ret__":
317 # print(name, "in mapping")
318 # if isinstance(name.ctx, ast.Store):
319 # return Name("tmp", ctx=Store())
320 # else:
321 # return copy.deepcopy(self.mapping[name.id])
322 # return name
324 def visit_Expr(self, expr):
325 if isinstance(expr.value, Name) and expr.value.id in self.mapping:
326 if self.mapping[expr.value.id] is not None:
327 return copy.deepcopy(self.mapping[expr.value.id])
328 return self.generic_visit(expr)