Issue #2527: Add a *globals* argument to timeit functions, in order to override the globals namespace in which the timed code is executed.

Patch by Ben Roberts.

5 files changed, 68 insertions, 19 deletions

diff --git a/Doc/library/timeit.rst b/Doc/library/timeit.rst
index 19b5e4e..dea1ba7 100644
--- a/Doc/library/timeit.rst
+++ b/Doc/library/timeit.rst
@@ -59,10 +59,15 @@ Python Interface
 The module defines three convenience functions and a public class:
 
 
-.. function:: timeit(stmt='pass', setup='pass', timer=<default timer>, number=1000000)
+.. function:: timeit(stmt='pass', setup='pass', timer=<default timer>, number=1000000, globals=None)
 
    Create a :class:`Timer` instance with the given statement, *setup* code and
    *timer* function and run its :meth:`.timeit` method with *number* executions.
+   The optional *globals* argument specifies a namespace in which to execute the
+   code.
+
+   .. versionchanged:: 3.5
+      The optional *globals* parameter was added.
 
    .. note::
 
@@ -71,12 +76,15 @@ The module defines three convenience functions and a public class:
         It will instead return the data specified by your return statement.
 
 
-.. function:: repeat(stmt='pass', setup='pass', timer=<default timer>, repeat=3, number=1000000)
+.. function:: repeat(stmt='pass', setup='pass', timer=<default timer>, repeat=3, number=1000000, globals=None)
 
    Create a :class:`Timer` instance with the given statement, *setup* code and
    *timer* function and run its :meth:`.repeat` method with the given *repeat*
-   count and *number* executions.
+   count and *number* executions.  The optional *globals* argument specifies a
+   namespace in which to execute the code.
 
+   .. versionchanged:: 3.5
+      The optional *globals* parameter was added.
 
 .. function:: default_timer()
 
@@ -86,7 +94,7 @@ The module defines three convenience functions and a public class:
       :func:`time.perf_counter` is now the default timer.
 
 
-.. class:: Timer(stmt='pass', setup='pass', timer=<timer function>)
+.. class:: Timer(stmt='pass', setup='pass', timer=<timer function>, globals=None)
 
    Class for timing execution speed of small code snippets.
 
@@ -94,7 +102,9 @@ The module defines three convenience functions and a public class:
    for setup, and a timer function.  Both statements default to ``'pass'``;
    the timer function is platform-dependent (see the module doc string).
    *stmt* and *setup* may also contain multiple statements separated by ``;``
-   or newlines, as long as they don't contain multi-line string literals.
+   or newlines, as long as they don't contain multi-line string literals.  The
+   statement will by default be executed within timeit's namespace; this behavior
+   can be controlled by passing a namespace to *globals*.
 
    To measure the execution time of the first statement, use the :meth:`.timeit`
    method.  The :meth:`.repeat` method is a convenience to call :meth:`.timeit`
@@ -105,6 +115,8 @@ The module defines three convenience functions and a public class:
    will then be executed by :meth:`.timeit`.  Note that the timing overhead is a
    little larger in this case because of the extra function calls.
 
+   .. versionchanged:: 3.5
+      The optional *globals* parameter was added.
 
    .. method:: Timer.timeit(number=1000000)
 
@@ -324,3 +336,17 @@ To give the :mod:`timeit` module access to functions you define, you can pass a
    if __name__ == '__main__':
        import timeit
        print(timeit.timeit("test()", setup="from __main__ import test"))
+
+Another option is to pass :func:`globals` to the  *globals* parameter, which will cause the code
+to be executed within your current global namespace.  This can be more convenient
+than individually specifying imports::
+
+   def f(x):
+       return x**2
+   def g(x):
+       return x**4
+   def h(x):
+       return x**8
+
+   import timeit
+   print(timeit.timeit('[func(42) for func in (f,g,h)]', globals=globals()))
diff --git a/Lib/test/test_timeit.py b/Lib/test/test_timeit.py
index 625fb8d..b3a96bb 100644
--- a/Lib/test/test_timeit.py
+++ b/Lib/test/test_timeit.py
@@ -86,9 +86,10 @@ class TestTimeit(unittest.TestCase):
     def fake_callable_stmt(self):
         self.fake_timer.inc()
 
-    def timeit(self, stmt, setup, number=None):
+    def timeit(self, stmt, setup, number=None, globals=None):
         self.fake_timer = FakeTimer()
-        t = timeit.Timer(stmt=stmt, setup=setup, timer=self.fake_timer)
+        t = timeit.Timer(stmt=stmt, setup=setup, timer=self.fake_timer,
+                globals=globals)
         kwargs = {}
         if number is None:
             number = DEFAULT_NUMBER
@@ -127,6 +128,17 @@ class TestTimeit(unittest.TestCase):
                 timer=FakeTimer())
         self.assertEqual(delta_time, 0)
 
+    def test_timeit_globals_args(self):
+        global _global_timer
+        _global_timer = FakeTimer()
+        t = timeit.Timer(stmt='_global_timer.inc()', timer=_global_timer)
+        self.assertRaises(NameError, t.timeit, number=3)
+        timeit.timeit(stmt='_global_timer.inc()', timer=_global_timer,
+                      globals=globals(), number=3)
+        local_timer = FakeTimer()
+        timeit.timeit(stmt='local_timer.inc()', timer=local_timer,
+                      globals=locals(), number=3)
+
     def repeat(self, stmt, setup, repeat=None, number=None):
         self.fake_timer = FakeTimer()
         t = timeit.Timer(stmt=stmt, setup=setup, timer=self.fake_timer)
diff --git a/Lib/timeit.py b/Lib/timeit.py
index ead2030..5971d37 100755
--- a/Lib/timeit.py
+++ b/Lib/timeit.py
@@ -60,6 +60,8 @@ default_number = 1000000
 default_repeat = 3
 default_timer = time.perf_counter
 
+_globals = globals
+
 # Don't change the indentation of the template; the reindent() calls
 # in Timer.__init__() depend on setup being indented 4 spaces and stmt
 # being indented 8 spaces.
@@ -94,7 +96,9 @@ class Timer:
     The constructor takes a statement to be timed, an additional
     statement used for setup, and a timer function.  Both statements
     default to 'pass'; the timer function is platform-dependent (see
-    module doc string).
+    module doc string).  If 'globals' is specified, the code will be
+    executed within that namespace (as opposed to inside timeit's
+    namespace).
 
     To measure the execution time of the first statement, use the
     timeit() method.  The repeat() method is a convenience to call
@@ -104,10 +108,12 @@ class Timer:
     multi-line string literals.
     """
 
-    def __init__(self, stmt="pass", setup="pass", timer=default_timer):
+    def __init__(self, stmt="pass", setup="pass", timer=default_timer,
+                 globals=None):
         """Constructor.  See class doc string."""
         self.timer = timer
-        ns = {}
+        local_ns = {}
+        global_ns = _globals() if globals is None else globals
         if isinstance(stmt, str):
             stmt = reindent(stmt, 8)
             if isinstance(setup, str):
@@ -115,19 +121,19 @@ class Timer:
                 src = template.format(stmt=stmt, setup=setup)
             elif callable(setup):
                 src = template.format(stmt=stmt, setup='_setup()')
-                ns['_setup'] = setup
+                local_ns['_setup'] = setup
             else:
                 raise ValueError("setup is neither a string nor callable")
-            self.src = src # Save for traceback display
+            self.src = src  # Save for traceback display
             code = compile(src, dummy_src_name, "exec")
-            exec(code, globals(), ns)
-            self.inner = ns["inner"]
+            exec(code, global_ns, local_ns)
+            self.inner = local_ns["inner"]
         elif callable(stmt):
             self.src = None
             if isinstance(setup, str):
                 _setup = setup
                 def setup():
-                    exec(_setup, globals(), ns)
+                    exec(_setup, global_ns, local_ns)
             elif not callable(setup):
                 raise ValueError("setup is neither a string nor callable")
             self.inner = _template_func(setup, stmt)
@@ -208,14 +214,14 @@ class Timer:
         return r
 
 def timeit(stmt="pass", setup="pass", timer=default_timer,
-           number=default_number):
+           number=default_number, globals=None):
     """Convenience function to create Timer object and call timeit method."""
-    return Timer(stmt, setup, timer).timeit(number)
+    return Timer(stmt, setup, timer, globals).timeit(number)
 
 def repeat(stmt="pass", setup="pass", timer=default_timer,
-           repeat=default_repeat, number=default_number):
+           repeat=default_repeat, number=default_number, globals=None):
     """Convenience function to create Timer object and call repeat method."""
-    return Timer(stmt, setup, timer).repeat(repeat, number)
+    return Timer(stmt, setup, timer, globals).repeat(repeat, number)
 
 def main(args=None, *, _wrap_timer=None):
     """Main program, used when run as a script.
diff --git a/Misc/ACKS b/Misc/ACKS
index bc96b93..a8ea6d7 100644
--- a/Misc/ACKS
+++ b/Misc/ACKS
@@ -1130,6 +1130,7 @@ Juan M. Bello Rivas
 Davide Rizzo
 Anthony Roach
 Carl Robben
+Ben Roberts
 Mark Roberts
 Andy Robinson
 Jim Robinson
diff --git a/Misc/NEWS b/Misc/NEWS
index da1fc32..727d529 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -124,6 +124,10 @@ Core and Builtins
 Library
 -------
 
+- Issue #2527: Add a *globals* argument to timeit functions, in order to
+  override the globals namespace in which the timed code is executed.
+  Patch by Ben Roberts.
+
 - Issue #22118: Switch urllib.parse to use RFC 3986 semantics for the
   resolution of relative URLs, rather than RFCs 1808 and 2396.
   Patch by Demian Brecht.