装饰多个 Python 函数

Question

假设我有很多对互补函数：

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)

本质上，对于每个函数 funcX，如果 flag == True 则它将 return 函数详细信息，但如果 flag == False 则它将调用其互补函数，comp_funcX，其中 return 是互补函数的详细信息。

不必在每个 funcX 中使用 if/else 语句，是否可以使用装饰器来清理或简化此代码？大致如下：

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

是这样的吗？

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}"

wrapper函数通常取*args和**kwargs，所以比较通用。

Answer 2

一种略有不同的方法是使用简单的包装函数而不是成熟的装饰器。将所有函数和补充函数打包到相应的列表中，并通过带有索引的包装器调用您的函数：

funcs = [func1, func2, ...]
comps = [comp_func1, comp_func2, ...]

def wrapper(i, arg1, arg2, flag):
    if flag:
        return funcs[i-1](arg1, arg2)
    else:
        return comps[i-1](arg1, arg2)

现在要调用 func1，您只需调用：

wrapper(1, arg1, arg2, flag)

这是基于第二个代码块中实现的功能。即每个函数 returns 其对应的数据 - 它们内部没有 flag 条件。

装饰多个 Python 函数

Decorating Multiple Python Functions

python

python-decorators