support top-level setup statements in Timer objects

Author: picnixz

Doc/library/timeit.rst

@@ -62,30 +62,44 @@ Python Interface
 The module defines three convenience functions and a public class:
 
 
-.. 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.
+.. function:: timeit(stmt='pass', setup='pass', timer=<default timer>,\
+                     number=1000000, globals=None,\
+                     *,\
+                     global_setup='pass')
+
+   Create a :class:`Timer` instance with the given statement, *setup* code,
+   *global_setup* global code and *timer* function and run its :meth:`.timeit`
+   method with the given *number* executions.
    The optional *globals* argument specifies a namespace in which to execute the
    code.
 
    .. versionchanged:: 3.5
       The optional *globals* parameter was added.
 
+   .. versionchanged:: next
+      The optional *global_setup* parameter was added.
+
 
-.. function:: repeat(stmt='pass', setup='pass', timer=<default timer>, repeat=5, number=1000000, globals=None)
+.. function:: repeat(stmt='pass', setup='pass', timer=<default timer>,\
+                     repeat=5, number=1000000, globals=None,\
+                     *,\
+                     global_setup='pass')
 
-   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.  The optional *globals* argument specifies a
-   namespace in which to execute the code.
+   Create a :class:`Timer` instance with the given statement, *setup* code,
+   *global_setup* global code and *timer* function and run its :meth:`.repeat`
+   method with the given *repeat* 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.
 
    .. versionchanged:: 3.7
       Default value of *repeat* changed from 3 to 5.
 
+   .. versionchanged:: next
+      The optional *global_setup* parameter was added.
+
 
 .. function:: default_timer()
 
@@ -96,36 +110,52 @@ 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>, globals=None)
+.. class:: Timer(stmt='pass', setup='pass', timer=<timer function>,\
+                 globals=None, *, global_setup='pass')
 
    Class for timing execution speed of small code snippets.
 
    The constructor takes a statement to be timed, an additional statement used
-   for setup, and a timer function.  Both statements default to ``'pass'``;
+   for setup and global setup, and a timer function.
+   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.  The
-   statement will by default be executed within timeit's namespace; this behavior
+
+   *stmt*, *setup* and *global_setup* may also contain multiple statements
+   separated by ``;`` 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*.
 
+   The distinction between *setup* and *global_setup* is that *setup* is
+   executed once per :meth:`.timeit` call while *global_setup* is executed
+   only at construction time. *global_setup* is typically useful to execute
+   top-level-only statements such as :ref:`future statements <future>` and
+   :ref:`wildcard imports <import>`.
+
    To measure the execution time of the first statement, use the :meth:`.timeit`
    method.  The :meth:`.repeat` and :meth:`.autorange` methods are convenience
    methods to call :meth:`.timeit` multiple times.
 
-   The execution time of *setup* is excluded from the overall timed execution run.
+   The execution time of *setup* or *global_setup* is excluded from the
+   overall timed execution run.
 
-   The *stmt* and *setup* parameters can also take objects that are callable
+   The *stmt*, *setup* and *global_setup* parameters can also take objects that are callable
    without arguments.  This will embed calls to them in a timer function that
    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.
 
+   .. versionchanged:: next
+      The optional *global_setup* parameter was added.
+
    .. method:: Timer.timeit(number=1000000)
 
-      Time *number* executions of the main statement.  This executes the setup
-      statement once, and then returns the time it takes to execute the main
+      Time *number* executions of the main statement.
+
+      This executes the setup statement once, and then returns the time it takes to execute the main
       statement a number of times.  The default timer returns seconds as a float.
       The argument is the number of times through the loop, defaulting to one
       million.  The main statement, the setup statement and the timer function
@@ -208,7 +238,7 @@ Command-Line Interface
 
 When called as a program from the command line, the following form is used::
 
-   python -m timeit [-n N] [-r N] [-u U] [-s S] [-p] [-v] [-h] [statement ...]
+   python -m timeit [-n N] [-r N] [-u U] [-s S] [-g S] [-p] [-v] [-h] [statement ...]
 
 Where the following options are understood:
 
@@ -224,8 +254,14 @@ Where the following options are understood:
 
 .. option:: -s S, --setup=S
 
+   statement to be executed once per timer repetition (default ``pass``)
+
+.. option:: -g S, --global-setup=S
+
    statement to be executed once initially (default ``pass``)
 
+   .. versionadded:: next
+
 .. option:: -p, --process
 
    measure process time, not wallclock time, using :func:`time.process_time`

Doc/whatsnew/3.15.rst

@@ -355,6 +355,21 @@ tarfile
   and :cve:`2025-4435`.)
 
 
+timeit
+------
+
+* :class:`~timeit.Timer` now support global setup statements which
+  are typically useful for top-level-only statements such as future
+  statements or wildcard imports.
+
+  Contrary to normal setup statements which are executed every time
+  a timer is repeated, global setup statements are executed exactly
+  once when constructing the ``Timer`` object and their side-effects
+  remain active for the entire object lifetime.
+
+  (Contributed by Bénédikt Tran in :gh:`137578`.)
+
+
 types
 ------
 

Lib/test/test_timeit.py

@@ -1,3 +1,4 @@
+import itertools
 import timeit
 import unittest
 import sys
@@ -23,8 +24,9 @@ class FakeTimer:
     BASE_TIME = 42.0
     def __init__(self, seconds_per_increment=1.0):
         self.count = 0
+        self.global_setup_calls = 0
         self.setup_calls = 0
-        self.seconds_per_increment=seconds_per_increment
+        self.seconds_per_increment = seconds_per_increment
         timeit._fake_timer = self
 
     def __call__(self):
@@ -33,6 +35,9 @@ def __call__(self):
     def inc(self):
         self.count += 1
 
+    def global_setup(self):
+        self.global_setup_calls += 1
+
     def setup(self):
         self.setup_calls += 1
 
@@ -81,6 +86,15 @@ def test_timer_invalid_stmt(self):
         self.assertRaises(SyntaxError, timeit.Timer,
                           setup='while False:\n  pass', stmt='  break')
 
+    def test_timer_invalid_global_setup(self):
+        self.assertRaises(ValueError, timeit.Timer, global_setup=None)
+        self.assertRaises(SyntaxError, timeit.Timer, global_setup='return')
+        self.assertRaises(SyntaxError, timeit.Timer, global_setup='yield')
+        self.assertRaises(SyntaxError, timeit.Timer, global_setup='yield from ()')
+        self.assertRaises(SyntaxError, timeit.Timer, global_setup='break')
+        self.assertRaises(SyntaxError, timeit.Timer, global_setup='continue')
+        self.assertRaises(SyntaxError, timeit.Timer, global_setup='  pass')
+
     def test_timer_invalid_setup(self):
         self.assertRaises(ValueError, timeit.Timer, setup=None)
         self.assertRaises(SyntaxError, timeit.Timer, setup='return')
@@ -91,30 +105,62 @@ def test_timer_invalid_setup(self):
         self.assertRaises(SyntaxError, timeit.Timer, setup='from timeit import *')
         self.assertRaises(SyntaxError, timeit.Timer, setup='  pass')
 
+    def test_timer_empty_global_setup(self):
+        timeit.Timer(global_setup='')
+        timeit.Timer(global_setup=' \n\t\f')
+        timeit.Timer(global_setup='# comment')
+
     def test_timer_empty_stmt(self):
         timeit.Timer(stmt='')
         timeit.Timer(stmt=' \n\t\f')
         timeit.Timer(stmt='# comment')
 
+    def test_timer_global_setup_future(self):
+        timer = timeit.Timer(stmt='print(f.__annotations__)',
+                             setup='def f(x: int): pass',
+                             global_setup='from __future__ import annotations')
+        with captured_stdout() as stdout:
+            timer.timeit(1)
+        self.assertEqual(stdout.getvalue().strip(), str({'x': 'int'}))
+
+    def test_timer_global_setup_side_effect(self):
+        timer = timeit.Timer(stmt='print(s)', global_setup='s = 1')
+        with captured_stdout() as stdout:
+            timer.timeit(5)
+        self.assertEqual(stdout.getvalue(), '1\n' * 5)
+
+    def test_timer_callable_global_setup_side_effect(self):
+        self.fake_timer = FakeTimer()
+        timer = timeit.Timer(global_setup=self.fake_callable_global_setup)
+        timer.timeit(5)
+        self.assertEqual(self.fake_timer.count, 0)
+        self.assertEqual(self.fake_timer.setup_calls, 0)
+        self.assertEqual(self.fake_timer.global_setup_calls, 1)
+
+    fake_global_setup = "import timeit\ntimeit._fake_timer.global_setup()"
     fake_setup = "import timeit\ntimeit._fake_timer.setup()"
     fake_stmt = "import timeit\ntimeit._fake_timer.inc()"
 
+    def fake_callable_global_setup(self):
+        self.fake_timer.global_setup()
+
     def fake_callable_setup(self):
         self.fake_timer.setup()
 
     def fake_callable_stmt(self):
         self.fake_timer.inc()
 
-    def timeit(self, stmt, setup, number=None, globals=None):
+    def timeit(self, stmt, setup, global_setup, number=None, globals=None):
         self.fake_timer = FakeTimer()
         t = timeit.Timer(stmt=stmt, setup=setup, timer=self.fake_timer,
-                globals=globals)
+                globals=globals, global_setup=global_setup)
         kwargs = {}
         if number is None:
             number = DEFAULT_NUMBER
         else:
             kwargs['number'] = number
         delta_time = t.timeit(**kwargs)
+        self.assertEqual(self.fake_timer.global_setup_calls, 1)
         self.assertEqual(self.fake_timer.setup_calls, 1)
         self.assertEqual(self.fake_timer.count, number)
         self.assertEqual(delta_time, number)
@@ -123,21 +169,20 @@ def timeit(self, stmt, setup, number=None, globals=None):
     #def test_timeit_default_iters(self):
     #    self.timeit(self.fake_stmt, self.fake_setup)
 
-    def test_timeit_zero_iters(self):
-        self.timeit(self.fake_stmt, self.fake_setup, number=0)
-
-    def test_timeit_few_iters(self):
-        self.timeit(self.fake_stmt, self.fake_setup, number=3)
-
-    def test_timeit_callable_stmt(self):
-        self.timeit(self.fake_callable_stmt, self.fake_setup, number=3)
-
-    def test_timeit_callable_setup(self):
-        self.timeit(self.fake_stmt, self.fake_callable_setup, number=3)
-
-    def test_timeit_callable_stmt_and_setup(self):
-        self.timeit(self.fake_callable_stmt,
-                self.fake_callable_setup, number=3)
+    def test_timeit_simple(self):
+        for number, (callable_stmt, callable_setup, callable_global_setup) in (
+            itertools.product([0, 3], itertools.product([True, False], repeat=3))
+        ):
+            with self.subTest(
+                number=number,
+                callable_stmt=callable_stmt,
+                callable_setup=callable_setup,
+                callable_global_setup=callable_global_setup,
+            ):
+                stmt = self.fake_stmt if callable_stmt else self.fake_callable_stmt
+                setup = self.fake_setup if callable_setup else self.fake_callable_setup
+                global_setup = self.fake_global_setup if callable_global_setup else self.fake_callable_global_setup
+                self.timeit(stmt, setup, global_setup, number=number)
 
     # Takes too long to run in debug build.
     #def test_timeit_function(self):
@@ -161,9 +206,10 @@ def test_timeit_globals_args(self):
         timeit.timeit(stmt='local_timer.inc()', timer=local_timer,
                       globals=locals(), number=3)
 
-    def repeat(self, stmt, setup, repeat=None, number=None):
+    def repeat(self, stmt, setup, global_setup, repeat=None, number=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,
+                         global_setup=global_setup)
         kwargs = {}
         if repeat is None:
             repeat = DEFAULT_REPEAT
@@ -174,6 +220,7 @@ def repeat(self, stmt, setup, repeat=None, number=None):
         else:
             kwargs['number'] = number
         delta_times = t.repeat(**kwargs)
+        self.assertEqual(self.fake_timer.global_setup_calls, 1)
         self.assertEqual(self.fake_timer.setup_calls, repeat)
         self.assertEqual(self.fake_timer.count, repeat * number)
         self.assertEqual(delta_times, repeat * [float(number)])
@@ -182,26 +229,23 @@ def repeat(self, stmt, setup, repeat=None, number=None):
     #def test_repeat_default(self):
     #    self.repeat(self.fake_stmt, self.fake_setup)
 
-    def test_repeat_zero_reps(self):
-        self.repeat(self.fake_stmt, self.fake_setup, repeat=0)
-
-    def test_repeat_zero_iters(self):
-        self.repeat(self.fake_stmt, self.fake_setup, number=0)
-
-    def test_repeat_few_reps_and_iters(self):
-        self.repeat(self.fake_stmt, self.fake_setup, repeat=3, number=5)
-
-    def test_repeat_callable_stmt(self):
-        self.repeat(self.fake_callable_stmt, self.fake_setup,
-                repeat=3, number=5)
-
-    def test_repeat_callable_setup(self):
-        self.repeat(self.fake_stmt, self.fake_callable_setup,
-                repeat=3, number=5)
-
-    def test_repeat_callable_stmt_and_setup(self):
-        self.repeat(self.fake_callable_stmt, self.fake_callable_setup,
-                repeat=3, number=5)
+    def test_repeat_simple(self):
+        for repeat, (callable_stmt, callable_setup, callable_global_setup) in (
+            itertools.product([0, 3], itertools.product([True, False], repeat=3))
+        ):
+            # do not use number = None for repeat > 0 as it would be too large
+            number = 5 if repeat > 0 else None
+            with self.subTest(
+                repeat=repeat,
+                number=number,
+                callable_stmt=callable_stmt,
+                callable_setup=callable_setup,
+                callable_global_setup=callable_global_setup,
+            ):
+                stmt = self.fake_stmt if callable_stmt else self.fake_callable_stmt
+                setup = self.fake_setup if callable_setup else self.fake_callable_setup
+                global_setup = self.fake_global_setup if callable_global_setup else self.fake_callable_global_setup
+                self.repeat(stmt, setup, global_setup, repeat=repeat, number=number)
 
     # Takes too long to run in debug build.
     #def test_repeat_function(self):
@@ -274,6 +318,18 @@ def test_main_fixed_iters(self):
         s = self.run_main(seconds_per_increment=2.0, switches=['-n35'])
         self.assertEqual(s, "35 loops, best of 5: 2 sec per loop\n")
 
+    def test_main_global_setup(self):
+        s = self.run_main(seconds_per_increment=2.0,
+                switches=['-n35', '-g', 'print("CustomSetup")'])
+        self.assertEqual(s, "CustomSetup\n" +
+                "35 loops, best of 5: 2 sec per loop\n")
+
+    def test_main_multiple_global_setups(self):
+        s = self.run_main(seconds_per_increment=2.0,
+                switches=['-n35', '-g', 'a = "CustomSetup"', '-g', 'print(a)'])
+        self.assertEqual(s, "CustomSetup\n" +
+                "35 loops, best of 5: 2 sec per loop\n")
+
     def test_main_setup(self):
         s = self.run_main(seconds_per_increment=2.0,
                 switches=['-n35', '-s', 'print("CustomSetup")'])