Decorating Multiple Python Functions

Question

Let's say I have many pairs of complementary functions:

def comp_func1(arg1, arg2):
    return f"comp_func1 {arg1} {arg2}"

def func1(arg1, arg2, flag=True):
    if flag:
        return f"func1 {arg1} {arg2}"
    else:
        return comp_func1(arg1, arg2)

def comp_func2(arg1, arg2):
    return f"comp_func2 {arg1} {arg2}"

def func2(arg1, arg2, flag=True):
    if flag:
        return f"func2 {arg1} {arg2}"
    else:
        return comp_func2(arg1, arg2)

Essentially, for each function funcX, if flag == True then it will return the function details but if flag == False then it will call its complementary function, comp_funcX, which returns the complementary function details.

Instead of having to use if/else statements in each funcX, is it possible to use decorators to clean up or simplify this code? Something along the lines of:

def decorator():
    """
    Check the `flag` and choose whether to call the complementary function
    """
    def wrapper():
        if flag:
            return func(arg1, arg2)
        else:
            return comp_func(args, arg2) 
    return wrapper

def comp_func1(arg1, arg2):
    return f"comp_func1 {arg1} {arg2}"

@decorator(comp_func1)
def func1(arg1, arg2, flag=True):
    return f"func1 {arg1} {arg2}"

def comp_func2(arg1, arg2):
    return f"comp_func2 {arg1} {arg2}"

@decorator(comp_func2)
def func2(arg1, arg2, flag=True):
    return f"func2 {arg1} {arg2}"

Answer 1

Something like this?

def some_decorator(comp_function):
    def decorator_with_argument(function):
        def wrapper(*args, **kwargs):
            if kwargs.get("flag", False):
                return function(*args, **kwargs)
            else:
                return comp_function(*args, **kwargs)
        return wrapper
    return decorator_with_argument


def comp_func1(arg1, arg2, *_, **__):
    return f"comp_func1 {arg1} {arg2}"


@some_decorator(comp_func1)
def func1(arg1, arg2, *_, **__):
    return f"func1 {arg1} {arg2}"

The wrapper function usually takes *args and **kwargs, so it is more general.