能不能用装饰器给里面的function多加一个参数？

Question

有没有一种（好的）方法可以使用装饰器向装饰后的（内部）function 添加另一个参数。我想我已经找到了一种可以完成其中一部分的方法，但看起来很老套并且至少有一个问题。

我的场景（我可能会以更好的方式解决，但只是想知道这里）是针对我运行的每日数据脚本，我希望某些函数具有“阶段”参数，这将为每个阶段做同样的事情. 例如，

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')

我发现自己在“做一些事情”之后用 4 个不同的函数写了完全相同的东西。 所以我想也许我可以使用装饰器基本上添加从if phase == 1到末尾的所有内容，并且只在每个 func 中包含它独有的do stuff 。

在一些测试中，我找到了一种方法来实现这一点，但我不确定这是否是最好的方法，而且文档字符串也变得一团糟……这是我目前所拥有的：

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)

有了这些，我就可以使用两个参数调用func并且它可以工作。 但是，当我调用help时，它表明它只需要param而不是phase 。

我还想将信息添加到任何 function 的文档字符串中，装饰器正在说明phase的作用。 在我们到达包装器之前，我能够通过在装饰器add_phase中放置add_to_doc...两行来实现这一点。 这导致help(func)给出文档字符串及其原始内容和添加内容。 但它仍然只显示一个参数。 此外，当我 hover 在 VSCode 中通过func时，那里的帮助文本只提供了 func 文档字符串，而不是添加了装饰器文档字符串信息。

那么，是否有一些我缺少的已经开发的方法来做到这一点？ 这种方法可能会失败吗？ 有没有办法更新文档字符串中参数中显示的内容？

最后，是的，我可以通过测试函数外部的phase并将它们全部放在一个 if 块中来更轻松地完成此特定任务，但这会很无聊。

Answer 1

只是不要 functools.wrap 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

然后你可以看到它

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

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

但实际上这似乎会很痛苦而且是个坏主意......