Search code examples
pythondecorator

How do I pass a function with arguments in another function?


I was learning decorators and wonder why my code isn't working.

So basically I have a simple function for calculating the nth Fibonacci number. The decorator function is to calculate the time it takes for the Fibonacci function to run.


def timetocomplete(f):
    def wrap_func(*args, **kwargs):
        start = time.time()
        f(*args)
        end = time.time()
        print(f"{f.__name__} took {end - start} seconds to complete")
    
    return wrap_func


@timetocomplete
def fib(a):
    if (a <= 1):
        return 1
    else:
        return (fib (a-2) + fib (a-1))
    
print(fib(5))

Commenting out the decorator, the fib function works fine but the using the decorator throws an error saying

TypeError: unsupported operand type(s) for +: 'NoneType' and 'NoneType'


Solution

  • You need to return the result of f in the wrap_func. This is because here:

    return (fib (a-2) + fib (a-1))

    If your function has been wrapped but the wrapper doesn't return anything, it will attempt to add Nonetype to Nonetype.

    def timetocomplete(f):
        def wrap_func(*args, **kwargs):
            start = time.time()
            res = f(*args)
            end = time.time()
            print(f"{f.__name__} took {end - start} seconds to complete")
            return res # <-- Here
    return wrap_func
    
    
    @timetocomplete
    def fib(a):
        if (a <= 1):
            return 1
        else:
            return (fib (a-2) + fib (a-1))
    
    print(fib(5))
    

    This will work, but it will show you the timetocomplete for each iteration of fib If you're looking to wrap the entire fib sequence in the timetocomplete function, something like this would probably work just fine

    def timetocomplete(f):
        def wrap_func(*args, **kwargs):
            start = time.time()
            res = f(*args)
            end = time.time()
            print(f"{f.__name__} took {end - start} seconds to complete")
            return res
    
        return wrap_func
    
    def fib(a):
        if (a <= 1):
            return 1
        else:
            return (fib (a-2) + fib (a-1))
    
    @timetocomplete
    def calculate_fib(a):
        return fib(a)
    
    print(calculate_fib(5))