Python decorator to time recursive functions

问题:

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?

回答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

回答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

回答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.

回答4:

Although it is not an overall solution to the problem of integrating recursion with decorators, for the problem of timing only, I have verified that the last element of the tuple of the times is the overall run time, as this is the time from the upper-most recursive call. Thus if you had

@timed
def rec():
    ...

to get the overall runtime given the original function definitions you could simply do

rec()[1]

Getting the result of the call, on the other hand, would then require recusing through the nested tuple:

def get(tup):
    if isinstance(tup, tuple):
        return get(tup[0])
    else:
        return tup

This might be too complicated to simply get the result of your function.