How to use python decorators with recursive functions?

Question

So I am making a program that shows how long each sorting algorithm takes to run on a specific list of unsorted numbers. I created a decorator function as follows

def timed(func):
    def wrapper(*args, **kwargs):
        start = datetime.now()
        val = func(*args, **kwargs)
        delta = datetime.now() - start
        print(delta.microseconds / 1000)
        return val

    return wrapper

Now this decorator works for all iterative functions as expected. However, the problem is with recursive functions. Take merge sort for example:

def _merge_sorted(l, r):
    i = j = 0
    temp = []

    while i < len(l) and j < len(r):
        if l[i] <= r[j]:
            temp.append(l[i])
            i += 1
        else:
            temp.append(r[j])
            j += 1

    temp.extend(l[i::])
    temp.extend(r[j::])

    return temp


@timed
def merge(arr):
    if len(arr) <= 1:
        return arr

    mid = len(arr) // 2

    left = merge(arr[:mid])
    right = merge(arr[mid:])

    return _merge_sorted(left, right)

When I call merge() on a list, instead of printing the time once, the program prints the time for what seems like every call in the call stack of this recursion. So the output for [5,4,3,2,1] is

I know the problem is caused by the decorator printing for each call in the call stack, but I can't find a way to fix this behavior to make it print the total time taken and make it work for both iterative and recursive functions.

user_423926 · Answer

Using Context manager

def factorial_recursive(n):
  if n == 1:
    return n
  else:
    return n*factorial_recursive(n-1)


def factorial_iterative(n):
  f = 1
  for i in range(n):
    f *= (i+1)
  return f

from contextlib import contextmanager

@contextmanager
def timeit():
    start = datetime.now()
    yield lambda: (datetime.now() - start).microseconds / 1000

with timeit() as t:
  print (factorial_recursive(10))
  print (f"Time take: {t()}")

with timeit() as t:
  print (factorial_iterative(10))
  print (f"Time take: {t()}")

with timeit() as t:
  print (merge([5,4,3,2,1]))
  print (f"Time take: {t()}")

output:

3628800
Time take: 0.078
3628800
Time take: 0.032
[1, 2, 3, 4, 5]
Time take: 0.101

OR

class Timeit():
  def __init__(self):
    self.start = datetime.now()  

  def __enter__(self):
    pass
     
  def __exit__(self, exc_type, exc_value, exc_traceback):
    delta = datetime.now() - self.start
    print (f"Time take: {delta.microseconds / 1000}")
 
with Timeit():
    print (factorial_recursive(10))

with Timeit():
    print (factorial_iterative(10))

How to use python decorators with recursive functions?

1 Answers