Search code examples
pythonpython-3.xpython-decorators

Can we call main function in its decorator function in Python?


I was trying to use python decorator over a function. I am taking input from the user and doubling it only if the entered input is 5. This needs to be achieved using decorator. So if any other number is entered by the user, the decorator should print the number is not 5, and it should start over and ask for the input again. This recursion should happen until user enters 5.

Here is what I have tried.

def check_five(func):
    def wrapper(*args,**kwargs):
        if args==(5,):
            return func(*args)
        else:
            print("Numbers is not 5.")
            return func(int(input("Enter a: ")))
    return wrapper


@check_five
def double(a):
    return a+a

print(double(int(input("Enter a: "))))

This is not a perfect solution, because, if user enters different number for the first time, it gives output as

Numbers is not 5. Enter a:

and asks for another input. But for second time user enters any number, obviously it runs the main function and produces the doubled output.

I tried another way, which works as expected.

def check_five(func):
    def wrapper(*args,**kwargs):
        if args==(5,):
            return func(*args)
        else:
            print("Numbers is not 5.")
            return double(int(input("Enter a: ")))
    return wrapper


@check_five
def double(a):
    return a+a

print(double(int(input("Enter a: "))))

In this solution I called the main function inside the decorator, you can find in else block of the decorator.

But I am not sure if this is a good practice. Because decorators are meant to be used over multiple different functions, so we can not call each main function inside the decorator.

Is there any other way to achieve this?

Note: I want to achieve this using decorator only.


Solution

  • And I mentioned in my comment, you'll keep your decorator generic and avoid the useless recursion with a while loop; for example:

    def check_five(func):
        def wrapper(*args,**kwargs):
            while args != (5,):
                args = (int(input("Enter a: ")), )
            return func(*args)
        return wrapper
    
    
    @check_five
    def double(a):
        return a+a
    
    print(double(int(input("Enter a: "))))
    

    Example:

    Enter a: >? 4
    Enter a: >? 4
    Enter a: >? 3
    Enter a: >? 5
    10
    

    Edit: note that though you allow any args, your checking whether args == (5,) will only allow 1-argument functions, so you'll probably want to modify that.