Python/Cython如何返回链式组合函数？替代if分支提升性能

你的需求是实现get_function，接收一个映射到运算函数的参数列表（0对应pow、1对应div、2对应mult），返回一个按顺序执行链式运算的函数。例如传入[0,1,2]时，返回的函数需计算a*b/(x**c)，对应运算链pow(x,c) → div(结果, b) → mult(结果, a)。

当前实现的问题在于，每次调用返回的fun时都会遍历参数并执行分支判断，高频调用下会产生不必要的性能损耗。下面提供两种优化方案，直接预组合运算逻辑，彻底消除分支判断。

方案1：预存操作序列（安全通用） #

先定义基础运算函数（假设你未定义）：

def div(a, b):
    return a / b

def mult(a, b):
    return a * b

优化后的get_function：

def get_function(function_params):
    # 建立操作码到函数的映射，仅初始化一次
    op_map = {0: pow, 1: div, 2: mult}
    
    # 预编译操作函数列表和对应的参数索引
    ops = []
    param_indices = []
    for idx, op_code in enumerate(function_params):
        ops.append(op_map[op_code])
        # 对应原逻辑中的params[-i]（i从1开始）
        param_indices.append(-(idx + 1))
    
    def fun(x, params):
        result = x
        for op, p_idx in zip(ops, param_indices):
            result = op(result, params[p_idx])
        return result
    return fun

原理 #

在get_function被调用时，就提前把所有操作对应的函数和参数索引存储起来。后续调用fun时，只需遍历预存列表执行运算，完全去掉了分支判断，性能比原实现提升显著。

方案2：动态生成无循环函数（性能最优） #

如果追求极致性能，可以直接动态生成嵌套运算的函数代码，连循环都省去：

def get_function(function_params):
    op_map = {0: "pow", 1: "div", 2: "mult"}
    # 拼接运算表达式字符串
    expr = "x"
    for idx, op_code in enumerate(function_params):
        func_name = op_map[op_code]
        expr = f"{func_name}({expr}, params[-(idx+1)])"
    
    # 动态生成函数
    exec(f"""def fun(x, params):
        return {expr}""")
    return fun

原理 #

当get_function([0,1,2])被调用时，会生成如下代码的函数：

def fun(x, params):
    return mult(div(pow(x, params[-1]), params[-2]), params[-3])

调用时直接执行嵌套运算，没有任何循环或分支，性能达到最优。注意：此方案使用exec，需确保function_params来自可信输入，避免代码注入风险。

测试验证 #

# 测试两种方案
func1 = get_function([0,1,2])
print(func1(2, [4, 2, 3]))  # 输出1.0（对应4*2/(2**3)=8/8=1.0）

func2 = get_function([0,1,2])
print(func2(2, [4, 2, 3]))  # 输出1.0

内容的提问来源于stack exchange，提问作者tjaqu787