How can I capture return value with Python timeit module? How can I capture return value with Python timeit module? python python

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