CODEWITHSUNDEEP
pythonrecursiondecoratorpython-decorators
wchargin
wchargin

Reputation: 16047

Value-modifying decorator on recursive Python function

I'm trying to make a timing decorator in Python. It works like this:

def function_timer(f):
  """
  Creates an function timer wrapper around a function.
  This is meant to be used as a decorator (see PEP 318).

  Input:
    f :: any function

  Output:
    A function that, when called, executes f with any arguments, while timing
    the execution. The return value of this function depends on the return
    value of f. If f returns
      None ::
        then this function returns None
      a dictionary d ::
        then this function returns a new dictionary with all the contents of
        d, but with an additional key 'time' equal to the execution time of
        the function. If the key 'time' is taken in the original dictionary,
        then 'time_' will be tried, then 'time__', etc., until a free key is
        found. The original dictionary will be shallow-copied, not modified.
      some non-dict value x:
        then this function returns a dictionary whose 'result' key gives x,
        and whose 'time' key gives the execution time of the function.
  """
  def timed_function(*args, **kwargs):
    start_time = time.time()
    print args, kwargs
    result = f(*args, **kwargs)
    end_time = time.time()
    total_time = end_time - start_time

    if result is None:
      return None
    elif type(result) >= dict:
      result = result.copy()
      key = 'time'
      while key in result:
        key += '_'
      result[key] = total_time
    else:
      return { 'result': result, 'time': total_time }

  return timed_function

This works for simple functions. But when I try it on a function like non-memoized Fibonacci

@function_timer
def fib(n):
  if n < 2:
    return 1
  else:
    return fib(n - 2) + fib(n - 1)

it fails, because as it recurses:

How can I get a decorator to only decorate the final return value of a function? I don't want to have to modify the fib function at all.

Upvotes: 2

Views: 218

Answers (3)

roippi
roippi

Reputation: 25974

There be hacks ahead. Proceed with caution.

Since cPython gives us access to frame introspection abilities, you can use that to try to figure out if a given call is recursive or not. Basically: if we're calling f, we can peek at the next frame up on the stack; if it's got the same __name__ as f, it's probably the result of a recursive call. Probably.

import sys

def function_timer(f):
   def timed_function(*args, **kwargs):
     if f.__name__ == sys._getframe(1).f_code.co_name:
         # recursive call!  Probably!
         return f(*args, **kwargs)
     # no changes to your function below here
     start_time = time.time()
     result = f(*args, **kwargs)
     end_time = time.time()
     total_time = end_time - start_time

     if result is None:
       return None
     elif type(result) >= dict:
       result = result.copy()
       key = 'time'
       while key in result:
         key += '_'
       result[key] = total_time
     else:
       return { 'result': result, 'time': total_time }

   return timed_function

So then:

@function_timer
def fib(n):
    if n < 2:
        return 1
    else:
        return fib(n - 2) + fib(n - 1)

fib(2)
Out[60]: {'result': 2, 'time': 2.86102294921875e-06}

fib(20)
Out[61]: {'result': 10946, 'time': 0.03364205360412598}

Disclaimer section: this is fragile. Don't do this in anything resembling production code - or if you must, come up with a way more vigorous test than if f.__name__ == sys._getframe(1).f_code.co_name:. And definitely don't expect this to be portable to all implementations of python.

All that said, diving into stack frames is basically the only way to go about figuring if a given function call is recursive. Or, well, there is technically another way by annotating the function's __dict__, but that's an order of magnitude more fragile, so this is as good as it gets. (Or simply don't use a decorator, as the other answers say)

Upvotes: 1

user2357112
user2357112

Reputation: 281683

timed_fib = function_timer(fib)

Trying to have fib provide one behavior to recursive calls and another to other users of the function is a recipe for pain. Just don't replace fib. Use another name for the timed function. If you want, you can rename fib to _fib and have the timed version be fib:

fib = function_timer(_fib)

Upvotes: 1

Martin Konecny
Martin Konecny

Reputation: 59681

If you use the function_timer function as a decorator, it will time each invocation of your fib function which you don't want. You probably want to time the cumulative time it takes for all the Fibonacci functions to execute.

As suggested by @minitech, you can simply wrap the fibonacci function via your function_timer function.

start_time = time.time()
result = fib(*args, **kwargs)
end_time = time.time()

The second line above will only return once all recursion is complete, and end_time - start_time will accurately reflect the execution time.

You may not realize this, but a decorator f wraps a function g like so: f(g(x)). In your case you don't want to wrap g with f for every call. You only want to it once, and so a decorator is not very useful here.

Upvotes: 2

Related Questions