From e6f50af5a16034355c4e1011e507a3bc289b3189 2012-09-26 13:22:50
From: Thomas Kluyver <takowl@gmail.com>
Date: 2012-09-26 13:22:50
Subject: [PATCH] Update %timeit magic to use AST transformations

---

diff --git a/IPython/core/magics/execution.py b/IPython/core/magics/execution.py
index 7d0b2d0..af73e1a 100644
--- a/IPython/core/magics/execution.py
+++ b/IPython/core/magics/execution.py
@@ -14,6 +14,7 @@
 
 # Stdlib
 import __builtin__ as builtin_mod
+import ast
 import bdb
 import os
 import sys
@@ -760,26 +761,54 @@ python-profiler package from non-free.""")
         # but is there a better way to achieve that the code stmt has access
         # to the shell namespace?
         transform  = self.shell.input_splitter.transform_cell
+        
         if cell is None:
             # called as line magic
-            setup = 'pass'
-            stmt = timeit.reindent(transform(stmt), 8)
-        else:
-            setup = timeit.reindent(transform(stmt), 4)
-            stmt = timeit.reindent(transform(cell), 8)
-
-        # From Python 3.3, this template uses new-style string formatting.
-        if sys.version_info >= (3, 3):
-            src = timeit.template.format(stmt=stmt, setup=setup)
+            ast_setup = ast.parse("pass")
+            ast_stmt = ast.parse(transform(stmt))
         else:
-            src = timeit.template % dict(stmt=stmt, setup=setup)
+            ast_setup = ast.parse(transform(stmt))
+            ast_stmt = ast.parse(transform(cell))
+        
+        ast_setup = self.shell.transform_ast(ast_setup)
+        ast_stmt = self.shell.transform_ast(ast_stmt)
+        
+        # This codestring is taken from timeit.template - we fill it in as an
+        # AST, so that we can apply our AST transformations to the user code
+        # without affecting the timing code.
+        timeit_ast_template = ast.parse('def inner(_it, _timer):\n'
+                                        '    setup\n'
+                                        '    _t0 = _timer()\n'
+                                        '    for _i in _it:\n'
+                                        '        stmt\n'
+                                        '    _t1 = _timer()\n'
+                                        '    return _t1 - _t0\n')
+        
+        class TimeitTemplateFiller(ast.NodeTransformer):
+            "This is quite tightly tied to the template definition above."
+            def visit_FunctionDef(self, node):
+                "Fill in the setup statement"
+                self.generic_visit(node)
+                if node.name == "inner":
+                    node.body[:1] = ast_setup.body
+                
+                return node
+            
+            def visit_For(self, node):
+                "Fill in the statement to be timed"
+                if getattr(getattr(node.body[0], 'value', None), 'id', None) == 'stmt':
+                    node.body = ast_stmt.body
+                return node
+        
+        timeit_ast = TimeitTemplateFiller().visit(timeit_ast_template)
+        timeit_ast = ast.fix_missing_locations(timeit_ast)
 
         # Track compilation time so it can be reported if too long
         # Minimum time above which compilation time will be reported
         tc_min = 0.1
 
         t0 = clock()
-        code = compile(src, "<magic-timeit>", "exec")
+        code = compile(timeit_ast, "<magic-timeit>", "exec")
         tc = clock()-t0
 
         ns = {}