Can you use a decorator to add an additional parameter to the inner function?

Viewed 73

Is there a (good) method for using a decorator to add another parameter to the decorated (inner) function. I think I've found a method that will do some of it, but seems hacky and has at least one problem.

My scenario (which I could likely solve in a better way, but just wondering here), is for a daily data script I run where I want some functions to have a parameter for 'phase' and this would do the same thing for each phase. For example,

def func(arg1,arg2,phase):
    # do some stuff specific to the function
    if phase == 1:
        pass
    elif phase == 2:
        # do something else, the same something in multiple different functions
    else:
        raise ValueError('phase needs to be 1 or 2')

I found myself writing the exact same thing in 4 different functions after the 'do some stuff'. So I thought that maybe I could use a decorator to basically add everything from the if phase == 1 to the end and just have the do stuff inside each func that is unique to it.

In some testing, I found a way to accomplish this, but I'm not sure if it is the best way, and also the docstrings get all messed up... Here's what I have so far:

import functools


def add_phase(func):
    add_to_doc = """\nhere is the add_phase docstring"""
    func.__doc__ += add_to_doc

    @functools.wraps(func)
    def wrapper(param, phase):
        """here is the wrapper docstring"""
        if phase == 1:
            print("got to this line")
        elif phase == 2:
            return func(param)
        else:
            raise ValueError
    return wrapper

@add_phase
def func(param):
    "here is the reg1 docstring"
    st = f"statement {param}"
    print(st)

func("param", phase=2)

With these, I am able to call func with both parameters and it works. However, when I call help on it it shows that it tells that it just takes param and not phase.

I also want to add the information to the docstring for any function that the decorator is on saying what phase does. I was able to accomplish that by putting the add_to_doc... two lines in the decorator add_phase before we get to the wrapper. That results in help(func) giving the docstring with its original plus the addition. But it still just shows the one param. Also, when I hover over func in VSCode, the help-text there just gives the func docstring, not with the decorator docstring info added.

So, is there some already-developed method to do this that I'm missing? Is this method likely going to break? Is there a way to update what shows in the params in the docstring?

Last, yes, I could probably more easily accomplish this specific task by testing phase outside the functions and having them all in an if block, but that would be boring.

1 Answers

just dont functools.wrap the function...

def add_phase(func):
    def wrapper(param, phase):
        if phase == 1:
            print("got to this line")
        elif phase == 2:
            return func(param)
        else:
            raise ValueError
    wrapper.__doc__ = f"Wraps: {func.__name__} and injects an arg\n{func.__doc__}"
    return wrapper

then you can see it

>>> help(func)
Help on function wrapper in module __main__:

wrapper(param, phase)
    Wraps: func and injects an arg
    here is the reg1 docstring

but really this seems like it will be fairly painful and a bad idea...

Related