How can I capture return value with Python timeit module?
For Python 3.5 you can override the value of timeit.template
timeit.template = """def inner(_it, _timer{init}): {setup} _t0 = _timer() for _i in _it: retval = {stmt} _t1 = _timer() return _t1 - _t0, retval"""
unutbu's answer works for python 3.4 but not 3.5 as the _template_func function appears to have been removed in 3.5
The problem boils down to timeit._template_func not returning the function's return value:
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
We can bend timeit
to our will with a bit of monkey-patching:
import timeitimport timedef _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: retval = _func() _t1 = _timer() return _t1 - _t0, retval return innertimeit._template_func = _template_funcdef foo(): time.sleep(1) return 42t = timeit.Timer(foo)print(t.timeit(number=1))
returns
(1.0010340213775635, 42)
The first value is the timeit result (in seconds), the second value is the function's return value.
Note that the monkey-patch above only affects the behavior of timeit
when a callable is passed timeit.Timer
. If you pass a string statement, then you'd have to (similarly) monkey-patch the timeit.template
string.
Funnily enough, I'm also doing machine-learning, and have a similar requirement ;-)
I solved it as follows, by writing a function, that:
- runs your function
- prints the running time, along with the name of your function
- returns the results
Let's say you want to time:
clf = RandomForest(train_input, train_output)
Then do:
clf = time_fn( RandomForest, train_input, train_output )
Stdout will show something like:
mymodule.RandomForest: 0.421609s
Code for time_fn:
import timedef time_fn( fn, *args, **kwargs ): start = time.clock() results = fn( *args, **kwargs ) end = time.clock() fn_name = fn.__module__ + "." + fn.__name__ print fn_name + ": " + str(end-start) + "s" return results