#11578: add unit tests for timeit module.

Patch by Michael Henry.

3 files changed, 313 insertions, 3 deletions

diff --git a/Lib/test/test_sundry.py b/Lib/test/test_sundry.py
index 4dacb9d..07802d6 100644
--- a/Lib/test/test_sundry.py
+++ b/Lib/test/test_sundry.py
@@ -54,7 +54,6 @@ class TestUntestedModules(unittest.TestCase):
             import py_compile
             import sndhdr
             import tabnanny
-            import timeit
             try:
                 import tty     # not available on Windows
             except ImportError:
diff --git a/Lib/test/test_timeit.py b/Lib/test/test_timeit.py
new file mode 100644
index 0000000..eb3b1a1
--- /dev/null
+++ b/Lib/test/test_timeit.py
@@ -0,0 +1,305 @@
+import timeit
+import unittest
+import sys
+import io
+import time
+from textwrap import dedent
+
+from test.support import run_unittest
+from test.support import captured_stdout
+from test.support import captured_stderr
+
+# timeit's default number of iterations.
+DEFAULT_NUMBER = 1000000
+
+# timeit's default number of repetitions.
+DEFAULT_REPEAT = 3
+
+# XXX: some tests are commented out that would improve the coverage but take a
+# long time to run because they test the default number of loops, which is
+# large.  The tests could be enabled if there was a way to override the default
+# number of loops during testing, but this would require changing the signature
+# of some functions that use the default as a default argument.
+
+class FakeTimer:
+    BASE_TIME = 42.0
+    def __init__(self, seconds_per_increment=1.0):
+        self.count = 0
+        self.setup_calls = 0
+        self.seconds_per_increment=seconds_per_increment
+        timeit._fake_timer = self
+
+    def __call__(self):
+        return self.BASE_TIME + self.count * self.seconds_per_increment
+
+    def inc(self):
+        self.count += 1
+
+    def setup(self):
+        self.setup_calls += 1
+
+    def wrap_timer(self, timer):
+        """Records 'timer' and returns self as callable timer."""
+        self.saved_timer = timer
+        return self
+
+class TestTimeit(unittest.TestCase):
+
+    def tearDown(self):
+        try:
+            del timeit._fake_timer
+        except AttributeError:
+            pass
+
+    def test_reindent_empty(self):
+        self.assertEqual(timeit.reindent("", 0), "")
+        self.assertEqual(timeit.reindent("", 4), "")
+
+    def test_reindent_single(self):
+        self.assertEqual(timeit.reindent("pass", 0), "pass")
+        self.assertEqual(timeit.reindent("pass", 4), "pass")
+
+    def test_reindent_multi_empty(self):
+        self.assertEqual(timeit.reindent("\n\n", 0), "\n\n")
+        self.assertEqual(timeit.reindent("\n\n", 4), "\n    \n    ")
+
+    def test_reindent_multi(self):
+        self.assertEqual(timeit.reindent(
+            "print()\npass\nbreak", 0),
+            "print()\npass\nbreak")
+        self.assertEqual(timeit.reindent(
+            "print()\npass\nbreak", 4),
+            "print()\n    pass\n    break")
+
+    def test_timer_invalid_stmt(self):
+        self.assertRaises(ValueError, timeit.Timer, stmt=None)
+
+    def test_timer_invalid_setup(self):
+        self.assertRaises(ValueError, timeit.Timer, setup=None)
+
+    fake_setup = "import timeit; timeit._fake_timer.setup()"
+    fake_stmt = "import timeit; timeit._fake_timer.inc()"
+
+    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):
+        self.fake_timer = FakeTimer()
+        t = timeit.Timer(stmt=stmt, setup=setup, timer=self.fake_timer)
+        kwargs = {}
+        if number is None:
+            number = DEFAULT_NUMBER
+        else:
+            kwargs['number'] = number
+        delta_time = t.timeit(**kwargs)
+        self.assertEqual(self.fake_timer.setup_calls, 1)
+        self.assertEqual(self.fake_timer.count, number)
+        self.assertEqual(delta_time, number)
+
+    # Takes too long to run in debug build.
+    #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_stmt_and_setup(self):
+        self.timeit(self.fake_callable_stmt,
+                self.fake_callable_setup, number=3)
+
+    # Takes too long to run in debug build.
+    #def test_timeit_function(self):
+    #    delta_time = timeit.timeit(self.fake_stmt, self.fake_setup,
+    #            timer=FakeTimer())
+    #    self.assertEqual(delta_time, DEFAULT_NUMBER)
+
+    def test_timeit_function_zero_iters(self):
+        delta_time = timeit.timeit(self.fake_stmt, self.fake_setup, number=0,
+                timer=FakeTimer())
+        self.assertEqual(delta_time, 0)
+
+    def repeat(self, stmt, setup, repeat=None, number=None):
+        self.fake_timer = FakeTimer()
+        t = timeit.Timer(stmt=stmt, setup=setup, timer=self.fake_timer)
+        kwargs = {}
+        if repeat is None:
+            repeat = DEFAULT_REPEAT
+        else:
+            kwargs['repeat'] = repeat
+        if number is None:
+            number = DEFAULT_NUMBER
+        else:
+            kwargs['number'] = number
+        delta_times = t.repeat(**kwargs)
+        self.assertEqual(self.fake_timer.setup_calls, repeat)
+        self.assertEqual(self.fake_timer.count, repeat * number)
+        self.assertEqual(delta_times, repeat * [float(number)])
+
+    # Takes too long to run in debug build.
+    #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_stmt_and_setup(self):
+        self.repeat(self.fake_callable_stmt, self.fake_callable_setup,
+                repeat=3, number=5)
+
+    # Takes too long to run in debug build.
+    #def test_repeat_function(self):
+    #    delta_times = timeit.repeat(self.fake_stmt, self.fake_setup,
+    #            timer=FakeTimer())
+    #    self.assertEqual(delta_times, DEFAULT_REPEAT * [float(DEFAULT_NUMBER)])
+
+    def test_repeat_function_zero_reps(self):
+        delta_times = timeit.repeat(self.fake_stmt, self.fake_setup, repeat=0,
+                timer=FakeTimer())
+        self.assertEqual(delta_times, [])
+
+    def test_repeat_function_zero_iters(self):
+        delta_times = timeit.repeat(self.fake_stmt, self.fake_setup, number=0,
+                timer=FakeTimer())
+        self.assertEqual(delta_times, DEFAULT_REPEAT * [0.0])
+
+    def assert_exc_string(self, exc_string, expected_exc_name):
+        exc_lines = exc_string.splitlines()
+        self.assertGreater(len(exc_lines), 2)
+        self.assertTrue(exc_lines[0].startswith('Traceback'))
+        self.assertTrue(exc_lines[-1].startswith(expected_exc_name))
+
+    def test_print_exc(self):
+        s = io.StringIO()
+        t = timeit.Timer("1/0")
+        try:
+            t.timeit()
+        except:
+            t.print_exc(s)
+        self.assert_exc_string(s.getvalue(), 'ZeroDivisionError')
+
+    MAIN_DEFAULT_OUTPUT = "10 loops, best of 3: 1 sec per loop\n"
+
+    def run_main(self, seconds_per_increment=1.0, switches=None, timer=None):
+        if timer is None:
+            timer = FakeTimer(seconds_per_increment=seconds_per_increment)
+        if switches is None:
+            args = []
+        else:
+            args = switches[:]
+        args.append(self.fake_stmt)
+        # timeit.main() modifies sys.path, so save and restore it.
+        orig_sys_path = sys.path[:]
+        with captured_stdout() as s:
+            timeit.main(args=args, _wrap_timer=timer.wrap_timer)
+        sys.path[:] = orig_sys_path[:]
+        return s.getvalue()
+
+    def test_main_bad_switch(self):
+        s = self.run_main(switches=['--bad-switch'])
+        self.assertEqual(s, dedent("""\
+            option --bad-switch not recognized
+            use -h/--help for command line help
+            """))
+
+    def test_main_seconds(self):
+        s = self.run_main(seconds_per_increment=5.5)
+        self.assertEqual(s, "10 loops, best of 3: 5.5 sec per loop\n")
+
+    def test_main_milliseconds(self):
+        s = self.run_main(seconds_per_increment=0.0055)
+        self.assertEqual(s, "100 loops, best of 3: 5.5 msec per loop\n")
+
+    def test_main_microseconds(self):
+        s = self.run_main(seconds_per_increment=0.0000025, switches=['-n100'])
+        self.assertEqual(s, "100 loops, best of 3: 2.5 usec per loop\n")
+
+    def test_main_fixed_iters(self):
+        s = self.run_main(seconds_per_increment=2.0, switches=['-n35'])
+        self.assertEqual(s, "35 loops, best of 3: 2 sec per loop\n")
+
+    def test_main_setup(self):
+        s = self.run_main(seconds_per_increment=2.0,
+                switches=['-n35', '-s', 'print("CustomSetup")'])
+        self.assertEqual(s, "CustomSetup\n" * 3 +
+                "35 loops, best of 3: 2 sec per loop\n")
+
+    def test_main_fixed_reps(self):
+        s = self.run_main(seconds_per_increment=60.0, switches=['-r9'])
+        self.assertEqual(s, "10 loops, best of 9: 60 sec per loop\n")
+
+    def test_main_negative_reps(self):
+        s = self.run_main(seconds_per_increment=60.0, switches=['-r-5'])
+        self.assertEqual(s, "10 loops, best of 1: 60 sec per loop\n")
+
+    def test_main_help(self):
+        s = self.run_main(switches=['-h'])
+        # Note: It's not clear that the trailing space was intended as part of
+        # the help text, but since it's there, check for it.
+        self.assertEqual(s, timeit.__doc__ + ' ')
+
+    def test_main_using_time(self):
+        fake_timer = FakeTimer()
+        s = self.run_main(switches=['-t'], timer=fake_timer)
+        self.assertEqual(s, self.MAIN_DEFAULT_OUTPUT)
+        self.assertIs(fake_timer.saved_timer, time.time)
+
+    def test_main_using_clock(self):
+        fake_timer = FakeTimer()
+        s = self.run_main(switches=['-c'], timer=fake_timer)
+        self.assertEqual(s, self.MAIN_DEFAULT_OUTPUT)
+        self.assertIs(fake_timer.saved_timer, time.clock)
+
+    def test_main_verbose(self):
+        s = self.run_main(switches=['-v'])
+        self.assertEqual(s, dedent("""\
+                10 loops -> 10 secs
+                raw times: 10 10 10
+                10 loops, best of 3: 1 sec per loop
+            """))
+
+    def test_main_very_verbose(self):
+        s = self.run_main(seconds_per_increment=0.000050, switches=['-vv'])
+        self.assertEqual(s, dedent("""\
+                10 loops -> 0.0005 secs
+                100 loops -> 0.005 secs
+                1000 loops -> 0.05 secs
+                10000 loops -> 0.5 secs
+                raw times: 0.5 0.5 0.5
+                10000 loops, best of 3: 50 usec per loop
+            """))
+
+    def test_main_exception(self):
+        with captured_stderr() as error_stringio:
+            s = self.run_main(switches=['1/0'])
+        self.assert_exc_string(error_stringio.getvalue(), 'ZeroDivisionError')
+
+    def test_main_exception_fixed_reps(self):
+        with captured_stderr() as error_stringio:
+            s = self.run_main(switches=['-n1', '1/0'])
+        self.assert_exc_string(error_stringio.getvalue(), 'ZeroDivisionError')
+
+
+def test_main():
+    run_unittest(TestTimeit)
+
+if __name__ == '__main__':
+    test_main()
diff --git a/Lib/timeit.py b/Lib/timeit.py
index 5c613fc..8e04645 100644
--- a/Lib/timeit.py
+++ b/Lib/timeit.py
@@ -232,10 +232,10 @@ def repeat(stmt="pass", setup="pass", timer=default_timer,
     """Convenience function to create Timer object and call repeat method."""
     return Timer(stmt, setup, timer).repeat(repeat, number)
 
-def main(args=None):
+def main(args=None, *, _wrap_timer=None):
     """Main program, used when run as a script.
 
-    The optional argument specifies the command line to be parsed,
+    The optional 'args' argument specifies the command line to be parsed,
     defaulting to sys.argv[1:].
 
     The return value is an exit code to be passed to sys.exit(); it
@@ -244,6 +244,10 @@ def main(args=None):
     When an exception happens during timing, a traceback is printed to
     stderr and the return value is 1.  Exceptions at other times
     (including the template compilation) are not caught.
+
+    '_wrap_timer' is an internal interface used for unit testing.  If it
+    is not None, it must be a callable that accepts a timer function
+    and returns another timer function (used for unit testing).
     """
     if args is None:
         args = sys.argv[1:]
@@ -289,6 +293,8 @@ def main(args=None):
     # directory)
     import os
     sys.path.insert(0, os.curdir)
+    if _wrap_timer is not None:
+        timer = _wrap_timer(timer)
     t = Timer(stmt, setup, timer)
     if number == 0:
         # determine number so that 0.2 <= total time < 2.0