diff -r 2f443a200b8c Lib/test/test_timeit.py --- a/Lib/test/test_timeit.py Wed Mar 16 17:33:24 2011 -0400 +++ b/Lib/test/test_timeit.py Wed Mar 16 18:07:56 2011 -0400 @@ -112,6 +112,9 @@ 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) @@ -161,6 +164,10 @@ 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) diff -r 2f443a200b8c Lib/timeit.py --- a/Lib/timeit.py Wed Mar 16 17:33:24 2011 -0400 +++ b/Lib/timeit.py Wed Mar 16 18:07:56 2011 -0400 @@ -91,17 +91,6 @@ """Helper to reindent a multi-line statement.""" return src.replace("\n", "\n" + " "*indent) -def _template_func(setup, func): - """Create a timer function. Used if the "statement" is a callable.""" - def inner(_it, _timer, _func=func): - setup() - _t0 = _timer() - for _i in _it: - _func() - _t1 = _timer() - return _t1 - _t0 - return inner - class Timer: """Class for timing execution speed of small code snippets. @@ -121,32 +110,27 @@ def __init__(self, stmt="pass", setup="pass", timer=default_timer): """Constructor. See class doc string.""" self.timer = timer + global_dict = globals().copy() + + # Store each callable parameter in params into global_dict and convert + # the parameter into a source code string to invoke the stored callable. + params = dict(stmt=stmt, setup=setup) + for name, obj in params.items(): + if not isinstance(obj, str): + if not hasattr(obj, '__call__'): + raise ValueError(name + " is neither a string nor callable") + callable_name = '_timeit_callable_%s' % name + global_dict[callable_name] = obj + params[name] = callable_name + '()' + + # Save source for traceback display. + self.src = template % { + 'stmt': reindent(params['stmt'], 8), + 'setup': reindent(params['setup'], 4)} + code = compile(self.src, dummy_src_name, "exec") ns = {} - if isinstance(stmt, str): - stmt = reindent(stmt, 8) - if isinstance(setup, str): - setup = reindent(setup, 4) - src = template % {'stmt': stmt, 'setup': setup} - elif hasattr(setup, '__call__'): - src = template % {'stmt': stmt, 'setup': '_setup()'} - ns['_setup'] = setup - else: - raise ValueError("setup is neither a string nor callable") - self.src = src # Save for traceback display - code = compile(src, dummy_src_name, "exec") - exec(code, globals(), ns) - self.inner = ns["inner"] - elif hasattr(stmt, '__call__'): - self.src = None - if isinstance(setup, str): - _setup = setup - def setup(): - exec(_setup, globals(), ns) - elif not hasattr(setup, '__call__'): - raise ValueError("setup is neither a string nor callable") - self.inner = _template_func(setup, stmt) - else: - raise ValueError("stmt is neither a string nor callable") + exec(code, global_dict, ns) + self.inner = ns["inner"] def print_exc(self, file=None): """Helper to print a traceback from the timed code.