1import timeit
2import unittest
3import sys
4import io
5import time
6from textwrap import dedent
7
8from test.support import captured_stdout
9from test.support import captured_stderr
10
11# timeit's default number of iterations.
12DEFAULT_NUMBER = 1000000
13
14# timeit's default number of repetitions.
15DEFAULT_REPEAT = 3
16
17# XXX: some tests are commented out that would improve the coverage but take a
18# long time to run because they test the default number of loops, which is
19# large.  The tests could be enabled if there was a way to override the default
20# number of loops during testing, but this would require changing the signature
21# of some functions that use the default as a default argument.
22
23class FakeTimer:
24    BASE_TIME = 42.0
25    def __init__(self, seconds_per_increment=1.0):
26        self.count = 0
27        self.setup_calls = 0
28        self.seconds_per_increment=seconds_per_increment
29        timeit._fake_timer = self
30
31    def __call__(self):
32        return self.BASE_TIME + self.count * self.seconds_per_increment
33
34    def inc(self):
35        self.count += 1
36
37    def setup(self):
38        self.setup_calls += 1
39
40    def wrap_timer(self, timer):
41        """Records 'timer' and returns self as callable timer."""
42        self.saved_timer = timer
43        return self
44
45class TestTimeit(unittest.TestCase):
46
47    def tearDown(self):
48        try:
49            del timeit._fake_timer
50        except AttributeError:
51            pass
52
53    def test_reindent_empty(self):
54        self.assertEqual(timeit.reindent("", 0), "")
55        self.assertEqual(timeit.reindent("", 4), "")
56
57    def test_reindent_single(self):
58        self.assertEqual(timeit.reindent("pass", 0), "pass")
59        self.assertEqual(timeit.reindent("pass", 4), "pass")
60
61    def test_reindent_multi_empty(self):
62        self.assertEqual(timeit.reindent("\n\n", 0), "\n\n")
63        self.assertEqual(timeit.reindent("\n\n", 4), "\n    \n    ")
64
65    def test_reindent_multi(self):
66        self.assertEqual(timeit.reindent(
67            "print()\npass\nbreak", 0),
68            "print()\npass\nbreak")
69        self.assertEqual(timeit.reindent(
70            "print()\npass\nbreak", 4),
71            "print()\n    pass\n    break")
72
73    def test_timer_invalid_stmt(self):
74        self.assertRaises(ValueError, timeit.Timer, stmt=None)
75        self.assertRaises(SyntaxError, timeit.Timer, stmt='return')
76        self.assertRaises(SyntaxError, timeit.Timer, stmt='yield')
77        self.assertRaises(SyntaxError, timeit.Timer, stmt='yield from ()')
78        self.assertRaises(SyntaxError, timeit.Timer, stmt='break')
79        self.assertRaises(SyntaxError, timeit.Timer, stmt='continue')
80        self.assertRaises(SyntaxError, timeit.Timer, stmt='from timeit import *')
81
82    def test_timer_invalid_setup(self):
83        self.assertRaises(ValueError, timeit.Timer, setup=None)
84        self.assertRaises(SyntaxError, timeit.Timer, setup='return')
85        self.assertRaises(SyntaxError, timeit.Timer, setup='yield')
86        self.assertRaises(SyntaxError, timeit.Timer, setup='yield from ()')
87        self.assertRaises(SyntaxError, timeit.Timer, setup='break')
88        self.assertRaises(SyntaxError, timeit.Timer, setup='continue')
89        self.assertRaises(SyntaxError, timeit.Timer, setup='from timeit import *')
90
91    fake_setup = "import timeit\ntimeit._fake_timer.setup()"
92    fake_stmt = "import timeit\ntimeit._fake_timer.inc()"
93
94    def fake_callable_setup(self):
95        self.fake_timer.setup()
96
97    def fake_callable_stmt(self):
98        self.fake_timer.inc()
99
100    def timeit(self, stmt, setup, number=None, globals=None):
101        self.fake_timer = FakeTimer()
102        t = timeit.Timer(stmt=stmt, setup=setup, timer=self.fake_timer,
103                globals=globals)
104        kwargs = {}
105        if number is None:
106            number = DEFAULT_NUMBER
107        else:
108            kwargs['number'] = number
109        delta_time = t.timeit(**kwargs)
110        self.assertEqual(self.fake_timer.setup_calls, 1)
111        self.assertEqual(self.fake_timer.count, number)
112        self.assertEqual(delta_time, number)
113
114    # Takes too long to run in debug build.
115    #def test_timeit_default_iters(self):
116    #    self.timeit(self.fake_stmt, self.fake_setup)
117
118    def test_timeit_zero_iters(self):
119        self.timeit(self.fake_stmt, self.fake_setup, number=0)
120
121    def test_timeit_few_iters(self):
122        self.timeit(self.fake_stmt, self.fake_setup, number=3)
123
124    def test_timeit_callable_stmt(self):
125        self.timeit(self.fake_callable_stmt, self.fake_setup, number=3)
126
127    def test_timeit_callable_setup(self):
128        self.timeit(self.fake_stmt, self.fake_callable_setup, number=3)
129
130    def test_timeit_callable_stmt_and_setup(self):
131        self.timeit(self.fake_callable_stmt,
132                self.fake_callable_setup, number=3)
133
134    # Takes too long to run in debug build.
135    #def test_timeit_function(self):
136    #    delta_time = timeit.timeit(self.fake_stmt, self.fake_setup,
137    #            timer=FakeTimer())
138    #    self.assertEqual(delta_time, DEFAULT_NUMBER)
139
140    def test_timeit_function_zero_iters(self):
141        delta_time = timeit.timeit(self.fake_stmt, self.fake_setup, number=0,
142                timer=FakeTimer())
143        self.assertEqual(delta_time, 0)
144
145    def test_timeit_globals_args(self):
146        global _global_timer
147        _global_timer = FakeTimer()
148        t = timeit.Timer(stmt='_global_timer.inc()', timer=_global_timer)
149        self.assertRaises(NameError, t.timeit, number=3)
150        timeit.timeit(stmt='_global_timer.inc()', timer=_global_timer,
151                      globals=globals(), number=3)
152        local_timer = FakeTimer()
153        timeit.timeit(stmt='local_timer.inc()', timer=local_timer,
154                      globals=locals(), number=3)
155
156    def repeat(self, stmt, setup, repeat=None, number=None):
157        self.fake_timer = FakeTimer()
158        t = timeit.Timer(stmt=stmt, setup=setup, timer=self.fake_timer)
159        kwargs = {}
160        if repeat is None:
161            repeat = DEFAULT_REPEAT
162        else:
163            kwargs['repeat'] = repeat
164        if number is None:
165            number = DEFAULT_NUMBER
166        else:
167            kwargs['number'] = number
168        delta_times = t.repeat(**kwargs)
169        self.assertEqual(self.fake_timer.setup_calls, repeat)
170        self.assertEqual(self.fake_timer.count, repeat * number)
171        self.assertEqual(delta_times, repeat * [float(number)])
172
173    # Takes too long to run in debug build.
174    #def test_repeat_default(self):
175    #    self.repeat(self.fake_stmt, self.fake_setup)
176
177    def test_repeat_zero_reps(self):
178        self.repeat(self.fake_stmt, self.fake_setup, repeat=0)
179
180    def test_repeat_zero_iters(self):
181        self.repeat(self.fake_stmt, self.fake_setup, number=0)
182
183    def test_repeat_few_reps_and_iters(self):
184        self.repeat(self.fake_stmt, self.fake_setup, repeat=3, number=5)
185
186    def test_repeat_callable_stmt(self):
187        self.repeat(self.fake_callable_stmt, self.fake_setup,
188                repeat=3, number=5)
189
190    def test_repeat_callable_setup(self):
191        self.repeat(self.fake_stmt, self.fake_callable_setup,
192                repeat=3, number=5)
193
194    def test_repeat_callable_stmt_and_setup(self):
195        self.repeat(self.fake_callable_stmt, self.fake_callable_setup,
196                repeat=3, number=5)
197
198    # Takes too long to run in debug build.
199    #def test_repeat_function(self):
200    #    delta_times = timeit.repeat(self.fake_stmt, self.fake_setup,
201    #            timer=FakeTimer())
202    #    self.assertEqual(delta_times, DEFAULT_REPEAT * [float(DEFAULT_NUMBER)])
203
204    def test_repeat_function_zero_reps(self):
205        delta_times = timeit.repeat(self.fake_stmt, self.fake_setup, repeat=0,
206                timer=FakeTimer())
207        self.assertEqual(delta_times, [])
208
209    def test_repeat_function_zero_iters(self):
210        delta_times = timeit.repeat(self.fake_stmt, self.fake_setup, number=0,
211                timer=FakeTimer())
212        self.assertEqual(delta_times, DEFAULT_REPEAT * [0.0])
213
214    def assert_exc_string(self, exc_string, expected_exc_name):
215        exc_lines = exc_string.splitlines()
216        self.assertGreater(len(exc_lines), 2)
217        self.assertTrue(exc_lines[0].startswith('Traceback'))
218        self.assertTrue(exc_lines[-1].startswith(expected_exc_name))
219
220    def test_print_exc(self):
221        s = io.StringIO()
222        t = timeit.Timer("1/0")
223        try:
224            t.timeit()
225        except:
226            t.print_exc(s)
227        self.assert_exc_string(s.getvalue(), 'ZeroDivisionError')
228
229    MAIN_DEFAULT_OUTPUT = "10 loops, best of 3: 1 sec per loop\n"
230
231    def run_main(self, seconds_per_increment=1.0, switches=None, timer=None):
232        if timer is None:
233            timer = FakeTimer(seconds_per_increment=seconds_per_increment)
234        if switches is None:
235            args = []
236        else:
237            args = switches[:]
238        args.append(self.fake_stmt)
239        # timeit.main() modifies sys.path, so save and restore it.
240        orig_sys_path = sys.path[:]
241        with captured_stdout() as s:
242            timeit.main(args=args, _wrap_timer=timer.wrap_timer)
243        sys.path[:] = orig_sys_path[:]
244        return s.getvalue()
245
246    def test_main_bad_switch(self):
247        s = self.run_main(switches=['--bad-switch'])
248        self.assertEqual(s, dedent("""\
249            option --bad-switch not recognized
250            use -h/--help for command line help
251            """))
252
253    def test_main_seconds(self):
254        s = self.run_main(seconds_per_increment=5.5)
255        self.assertEqual(s, "10 loops, best of 3: 5.5 sec per loop\n")
256
257    def test_main_milliseconds(self):
258        s = self.run_main(seconds_per_increment=0.0055)
259        self.assertEqual(s, "100 loops, best of 3: 5.5 msec per loop\n")
260
261    def test_main_microseconds(self):
262        s = self.run_main(seconds_per_increment=0.0000025, switches=['-n100'])
263        self.assertEqual(s, "100 loops, best of 3: 2.5 usec per loop\n")
264
265    def test_main_fixed_iters(self):
266        s = self.run_main(seconds_per_increment=2.0, switches=['-n35'])
267        self.assertEqual(s, "35 loops, best of 3: 2 sec per loop\n")
268
269    def test_main_setup(self):
270        s = self.run_main(seconds_per_increment=2.0,
271                switches=['-n35', '-s', 'print("CustomSetup")'])
272        self.assertEqual(s, "CustomSetup\n" * 3 +
273                "35 loops, best of 3: 2 sec per loop\n")
274
275    def test_main_multiple_setups(self):
276        s = self.run_main(seconds_per_increment=2.0,
277                switches=['-n35', '-s', 'a = "CustomSetup"', '-s', 'print(a)'])
278        self.assertEqual(s, "CustomSetup\n" * 3 +
279                "35 loops, best of 3: 2 sec per loop\n")
280
281    def test_main_fixed_reps(self):
282        s = self.run_main(seconds_per_increment=60.0, switches=['-r9'])
283        self.assertEqual(s, "10 loops, best of 9: 60 sec per loop\n")
284
285    def test_main_negative_reps(self):
286        s = self.run_main(seconds_per_increment=60.0, switches=['-r-5'])
287        self.assertEqual(s, "10 loops, best of 1: 60 sec per loop\n")
288
289    @unittest.skipIf(sys.flags.optimize >= 2, "need __doc__")
290    def test_main_help(self):
291        s = self.run_main(switches=['-h'])
292        # Note: It's not clear that the trailing space was intended as part of
293        # the help text, but since it's there, check for it.
294        self.assertEqual(s, timeit.__doc__ + ' ')
295
296    def test_main_using_time(self):
297        fake_timer = FakeTimer()
298        s = self.run_main(switches=['-t'], timer=fake_timer)
299        self.assertEqual(s, self.MAIN_DEFAULT_OUTPUT)
300        self.assertIs(fake_timer.saved_timer, time.time)
301
302    def test_main_using_clock(self):
303        fake_timer = FakeTimer()
304        s = self.run_main(switches=['-c'], timer=fake_timer)
305        self.assertEqual(s, self.MAIN_DEFAULT_OUTPUT)
306        self.assertIs(fake_timer.saved_timer, time.clock)
307
308    def test_main_verbose(self):
309        s = self.run_main(switches=['-v'])
310        self.assertEqual(s, dedent("""\
311                10 loops -> 10 secs
312                raw times: 10 10 10
313                10 loops, best of 3: 1 sec per loop
314            """))
315
316    def test_main_very_verbose(self):
317        s = self.run_main(seconds_per_increment=0.000050, switches=['-vv'])
318        self.assertEqual(s, dedent("""\
319                10 loops -> 0.0005 secs
320                100 loops -> 0.005 secs
321                1000 loops -> 0.05 secs
322                10000 loops -> 0.5 secs
323                raw times: 0.5 0.5 0.5
324                10000 loops, best of 3: 50 usec per loop
325            """))
326
327    def test_main_with_time_unit(self):
328        unit_sec = self.run_main(seconds_per_increment=0.002,
329                switches=['-u', 'sec'])
330        self.assertEqual(unit_sec,
331                "1000 loops, best of 3: 0.002 sec per loop\n")
332        unit_msec = self.run_main(seconds_per_increment=0.002,
333                switches=['-u', 'msec'])
334        self.assertEqual(unit_msec,
335                "1000 loops, best of 3: 2 msec per loop\n")
336        unit_usec = self.run_main(seconds_per_increment=0.002,
337                switches=['-u', 'usec'])
338        self.assertEqual(unit_usec,
339                "1000 loops, best of 3: 2e+03 usec per loop\n")
340        # Test invalid unit input
341        with captured_stderr() as error_stringio:
342            invalid = self.run_main(seconds_per_increment=0.002,
343                    switches=['-u', 'parsec'])
344        self.assertEqual(error_stringio.getvalue(),
345                    "Unrecognized unit. Please select usec, msec, or sec.\n")
346
347    def test_main_exception(self):
348        with captured_stderr() as error_stringio:
349            s = self.run_main(switches=['1/0'])
350        self.assert_exc_string(error_stringio.getvalue(), 'ZeroDivisionError')
351
352    def test_main_exception_fixed_reps(self):
353        with captured_stderr() as error_stringio:
354            s = self.run_main(switches=['-n1', '1/0'])
355        self.assert_exc_string(error_stringio.getvalue(), 'ZeroDivisionError')
356
357    def autorange(self, callback=None):
358        timer = FakeTimer(seconds_per_increment=0.001)
359        t = timeit.Timer(stmt=self.fake_stmt, setup=self.fake_setup, timer=timer)
360        return t.autorange(callback)
361
362    def test_autorange(self):
363        num_loops, time_taken = self.autorange()
364        self.assertEqual(num_loops, 1000)
365        self.assertEqual(time_taken, 1.0)
366
367    def test_autorange_with_callback(self):
368        def callback(a, b):
369            print("{} {:.3f}".format(a, b))
370        with captured_stdout() as s:
371            num_loops, time_taken = self.autorange(callback)
372        self.assertEqual(num_loops, 1000)
373        self.assertEqual(time_taken, 1.0)
374        expected = ('10 0.010\n'
375                    '100 0.100\n'
376                    '1000 1.000\n')
377        self.assertEqual(s.getvalue(), expected)
378
379
380if __name__ == '__main__':
381    unittest.main()
382