Search code examples
pythonrecursiondecorator

Tracking recursion depth with Python decorators


I am trying to write a decorator that tracks the recursion depth of a recursive function in Python.

Take, for example, a recursive function such as


def fib(n):
    if n == 0:
        return 0
    if n == 1:
        return 1
    else:
        return fib(n-1) + fib(n-2)

Usually, to track recursion depth you could write something like so


    def fib(n, level=0):
        print(level)
        if n == 0:
            return 0
        if n == 1:
            return 1
        else:
            return fib(n-1, level=level+1) + fib(n-2, level=level+1)

However, after fiddling around with decorators for a while and a lot of googling, I'm not sure this is even possible.

I have tried things like


def visualise(func):
    def modify(*args, **kwargs):
        kwargs["count"] += 1
        print(kwargs["count"])
        return func(*args, **kwargs)

    return modify

@visualise
def fib(n):
    ...

fib(4, count=0)

But it complains that count is an unexpected keyword argument, which I don't really understand because I was under the impression that the wrapper modify replaces all occurrences of fib but I guess not?

Pointers would be greatly appreciated.


Solution

  • You can specify level variable for decorator and then use it in function calls like this:

    def visualise(func):
         visualise.level = 0
         def wrapper(*args, **kwargs):
             print("In:", visualise.level)
             visualise.level += 1
             result = func(*args, **kwargs)
             visualise.level -= 1
             print("Out:", visualise.level)
             return result
         return wrapper