可以使用装饰器向内部函数添加额外的参数吗?



是否有一种(好的)方法可以使用decorator向已装饰的(内部)函数添加另一个参数?我想我已经找到了一种方法,可以做一些事情,但看起来很粗糙,至少有一个问题。

我的场景(我可能会以更好的方式解决,但只是想知道这里),是我运行的日常数据脚本,我想要一些函数有一个参数"阶段",这将为每个阶段做同样的事情。例如,

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到末尾的所有内容,并且只将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

我还想将信息添加到docstring中,用于装饰器所说的phase所做的任何函数。在我们进入包装器之前,我可以通过在装饰器add_phase中放入add_to_doc...两行来实现这一点。这将导致help(func)给出docstring的原始值和附加值。但它仍然只显示了一个参数。此外,当我在VSCode中悬停在func上时,那里的帮助文本只是给出了函数docstring,而不是添加了装饰器docstring info。

那么,是否有一些已经开发的方法来做到这一点,我错过了?这个方法可能会失败吗?是否有一种方法来更新在文档字符串的参数中显示的内容?

最后,是的,我可以更容易地完成这个特定的任务,通过在函数之外测试phase,并将它们都放在if块中,但那将是无聊的。

不要使用functools。包装函数…

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 argn{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

但实际上这看起来会相当痛苦,而且是个坏主意…

最新更新