Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.9/dist-packages/IPython/core/magics/ast_mod.py: 28%
90 statements
« prev ^ index » next coverage.py v7.3.1, created at 2023-09-25 06:05 +0000
« prev ^ index » next coverage.py v7.3.1, created at 2023-09-25 06:05 +0000
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 mesurement.
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 NodeTransformer, Store, Load, Name, Expr, Assign, Module
182import ast
183import copy
185from typing import Dict, Optional
188mangle_all = lambda name: False if name in ("__ret__", "__code__") else True
191class Mangler(NodeTransformer):
192 """
193 Mangle given names in and ast tree to make sure they do not conflict with
194 user code.
195 """
197 enabled: bool = True
198 debug: bool = False
200 def log(self, *args, **kwargs):
201 if self.debug:
202 print(*args, **kwargs)
204 def __init__(self, predicate=None):
205 if predicate is None:
206 predicate = lambda name: name.startswith("___")
207 self.predicate = predicate
209 def visit_Name(self, node):
210 if self.predicate(node.id):
211 self.log("Mangling", node.id)
212 # Once in the ast we do not need
213 # names to be valid identifiers.
214 node.id = "mangle-" + node.id
215 else:
216 self.log("Not mangling", node.id)
217 return node
219 def visit_FunctionDef(self, node):
220 if self.predicate(node.name):
221 self.log("Mangling", node.name)
222 node.name = "mangle-" + node.name
223 else:
224 self.log("Not mangling", node.name)
226 for arg in node.args.args:
227 if self.predicate(arg.arg):
228 self.log("Mangling function arg", arg.arg)
229 arg.arg = "mangle-" + arg.arg
230 else:
231 self.log("Not mangling function arg", arg.arg)
232 return self.generic_visit(node)
234 def visit_ImportFrom(self, node):
235 return self._visit_Import_and_ImportFrom(node)
237 def visit_Import(self, node):
238 return self._visit_Import_and_ImportFrom(node)
240 def _visit_Import_and_ImportFrom(self, node):
241 for alias in node.names:
242 asname = alias.name if alias.asname is None else alias.asname
243 if self.predicate(asname):
244 new_name: str = "mangle-" + asname
245 self.log("Mangling Alias", new_name)
246 alias.asname = new_name
247 else:
248 self.log("Not mangling Alias", alias.asname)
249 return node
252class ReplaceCodeTransformer(NodeTransformer):
253 enabled: bool = True
254 debug: bool = False
255 mangler: Mangler
257 def __init__(
258 self, template: Module, mapping: Optional[Dict] = None, mangling_predicate=None
259 ):
260 assert isinstance(mapping, (dict, type(None)))
261 assert isinstance(mangling_predicate, (type(None), type(lambda: None)))
262 assert isinstance(template, ast.Module)
263 self.template = template
264 self.mangler = Mangler(predicate=mangling_predicate)
265 if mapping is None:
266 mapping = {}
267 self.mapping = mapping
269 @classmethod
270 def from_string(
271 cls, template: str, mapping: Optional[Dict] = None, mangling_predicate=None
272 ):
273 return cls(
274 ast.parse(template), mapping=mapping, mangling_predicate=mangling_predicate
275 )
277 def visit_Module(self, code):
278 if not self.enabled:
279 return code
280 # if not isinstance(code, ast.Module):
281 # recursively called...
282 # return generic_visit(self, code)
283 last = code.body[-1]
284 if isinstance(last, Expr):
285 code.body.pop()
286 code.body.append(Assign([Name("ret-tmp", ctx=Store())], value=last.value))
287 ast.fix_missing_locations(code)
288 ret = Expr(value=Name("ret-tmp", ctx=Load()))
289 ret = ast.fix_missing_locations(ret)
290 self.mapping["__ret__"] = ret
291 else:
292 self.mapping["__ret__"] = ast.parse("None").body[0]
293 self.mapping["__code__"] = code.body
294 tpl = ast.fix_missing_locations(self.template)
296 tx = copy.deepcopy(tpl)
297 tx = self.mangler.visit(tx)
298 node = self.generic_visit(tx)
299 node_2 = ast.fix_missing_locations(node)
300 if self.debug:
301 print("---- Transformed code ----")
302 print(ast.unparse(node_2))
303 print("---- ---------------- ----")
304 return node_2
306 # this does not work as the name might be in a list and one might want to extend the list.
307 # def visit_Name(self, name):
308 # if name.id in self.mapping and name.id == "__ret__":
309 # print(name, "in mapping")
310 # if isinstance(name.ctx, ast.Store):
311 # return Name("tmp", ctx=Store())
312 # else:
313 # return copy.deepcopy(self.mapping[name.id])
314 # return name
316 def visit_Expr(self, expr):
317 if isinstance(expr.value, Name) and expr.value.id in self.mapping:
318 if self.mapping[expr.value.id] is not None:
319 return copy.deepcopy(self.mapping[expr.value.id])
320 return self.generic_visit(expr)