Python decorator to time recursive functions

Question

I have a simple decorator to track the runtime of a function call:

def timed(f):
    def caller(*args):
        start = time.time()
        res = f(*args)
        end = time.time()
        return res, end - start
    return caller

This can be used as follows, and returns a tuple of the function result and the execution time.

@timed
def test(n):
    for _ in range(n):
        pass
    return 0

print(test(900)) # prints (0, 2.69e-05)

Simple enough. But now I want to apply this to recursive functions. Applying the above wrapper to a recursive function results in nested tuples with the times of each recursive call, as is expected.

@timed
def rec(n):
    if n:
        return rec(n - 1)
    else:
        return 0

print(rec(3)) # Prints ((((0, 1.90e-06), 8.10e-06), 1.28e-05), 1.90e-05)

What's an elegant way to write the decorator so that it handles recursion properly? Obviously, you could wrap the call if a timed function:

@timed
def wrapper():
    return rec(3)

This will give a tuple of the result and the time, but I want all of it to be handled by the decorator so that the caller does not need to worry about defining a new function for every call. Ideas?

Answer 1

The problem here isn't really the decorator. The problem is that rec needs rec to be a function that behaves one way, but you want rec to be a function that behaves differently. There's no clean way to reconcile that with a single rec function.

The cleanest option is to stop requiring rec to be two things at once. Instead of using decorator notation, assign timed(rec) to a different name:

def rec(n):
    ...

timed_rec = timed(rec)

If you don't want two names, then rec needs to be written to understand the actual value that the decorated rec will return. For example,

@timed
def rec(n):
    if n:
        val, runtime = rec(n-1)
        return val
    else:
        return 0

Answer 2

I prefer the other answers so far (particularly user2357112's answer), but you can also make a class-based decorator that detects whether the function has been activated, and if so, bypasses the timing:

import time

class fancy_timed(object):
    def __init__(self, f):
        self.f = f
        self.active = False

    def __call__(self, *args):
        if self.active:
            return self.f(*args)
        start = time.time()
        self.active = True
        res = self.f(*args)
        end = time.time()
        self.active = False
        return res, end - start


@fancy_timed
def rec(n):
    if n:
        time.sleep(0.01)
        return rec(n - 1)
    else:
        return 0
print(rec(3))

(class written with (object) so that this is compatible with py2k and py3k).

Note that to really work properly, the outermost call should use try and finally. Here's the fancied up fancy version of __call__:

def __call__(self, *args):
    if self.active:
        return self.f(*args)
    try:
        start = time.time()
        self.active = True
        res = self.f(*args)
        end = time.time()
        return res, end - start
    finally:
        self.active = False

Answer 3

You could structure your timer in a different way by *ahem* abusing the contextmanager and function attribute a little...

from contextlib import contextmanager
import time

@contextmanager
def timed(func):
    timed.start = time.time()
    try:
        yield func
    finally:
        timed.duration = time.time() - timed.start

def test(n):
    for _ in range(n):
        pass
    return n

def rec(n):
    if n:
        time.sleep(0.05) # extra delay to notice the difference
        return rec(n - 1)
    else:
        return n

with timed(rec) as r:
    print(t(10))
    print(t(20))

print(timed.duration)

with timed(test) as t:
    print(t(555555))
    print(t(666666))

print(timed.duration)

Results:

# recursive
0
0
1.5130000114440918

# non-recursive
555555
666666
0.053999900817871094

If this is deemed a bad hack I'll gladly accept your criticism.