Python decorator with a function that receives a parameter

Question

Decorator is not working as intended

I was trying to understand more about decorators, so I searched and got many videos that made this exact function:

def tictoc(func):
    def wrapper():
        t1 = time.time()
        return func()
        t2 = time.time()-t1
        print(f'Took {t2} seconds to run.')
    return wrapper

So I made a factorial function, that takes a parameter (the number you want to see the factorial)

@tictoc
def fat(num):
    f = 1
    for x in range(1, num+1):
        f *= x
    return f

print(fat(5))

and then I ran it

but I got this error: TypeError: tictoc..wrapper() takes 0 positional arguments but 1 was given

Someone can please explain me what I am doing wrong?

Answer 1

Problem #1: The decorator does not allow passing of the argument as you are calling your func at return without arguments

def tictoc(func):
    def wrapper():
        t1 = time.time()
        return func() #<—- here is the 🐛
        t2 = time.time()-t1
        print(f'Took {t2} seconds to run.')
    return wrapper

We can correct that:

def tictoc(func):
    def wrapper(*args, **kwargs):
        t1 = time.time()
        return func(*args, **kwargs) 
        t2 = time.time()-t1 # <— this will not be evaluated as we have returned 
        print(f'Took {t2} seconds to run.') # <- this too
    return wrapper

Problem #2: Return statement will return. t2 and print as pointed in the comment above will not be evaluated. We can fix that:

def tictoc(func):
    def wrapper(*args, **kwargs):
        t1 = time.time()
        results = func(*args, **kwargs) 
        t2 = time.time()-t1 
        print(f'Took {t2} seconds to run.')
        return results 
    return wrapper

This should work now. We can add one last tweak to preserve the name of our function

As it is, the decorator will change the name of your function.

@tictoc
def fat(num):
    f = 1
    for x in range(1, num+1):
        f *= x
    return f

print(fat.__name__)
#outputs wrapper

To fix that

from functools import wraps


def tictoc(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        t1 = time.time()
        results = func(*args, **kwargs) 
        t2 = time.time()-t1 
        print(f'Took {t2} seconds to run.')
        return results 
    return wrapper

The output of print(fat.__name__) is not correct fat.

Another way to write your decorator could be using a class which is easier to read.

class tictoc:

    def __init__(self, func):
        self.func = func
        self.__name__ = func.__name__

    def __call__(self, *args, **kwargs):

        t1 = time.time()

        results = self.func(*args, **kwargs)
        t2 = time.time() - t1
        print(f'Took {t2} seconds to run.')
        return results

This should work as above but gives less freedom. For example adding parameters to your decorator. With function we can do that with three nested functions

def tictoc(loops=1):
    def inner(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            
            t2_ = 0
            for loop in range(1, loops+1):
                print(f'{loop=}')
                t1 = time.time()
                results = func(*args, **kwargs)
                t2 = time.time() - t1
                t2_ += t2
                
                print(f'Took {t2} seconds to run')
            if loops > 1:
                print(f'Took average {t2_/loops} seconds to run')
                
            return results
    
        return wrapper
    return inner




@tictoc(3)
def fat(num):
    f = 1
    for x in range(1, num + 1):
        f *= x
    return f

print(fat(5))

Still pondering how to do it with class ;)