如何替代函数中的if-elif实现多计算工具函数的选择？

你可以用字典映射的方式重构这段代码，既消除重复逻辑，也让后续新增计算工具的扩展更简单。

重构后的代码示例 #

def average_day_range(price_df: PandasDataFrame, n: int=14, calculation_tool: int=0):
    '''
    0 - SMA, 1 - EMA, 2 - WMA, 3 - EWMA, 4 - VWMA,  5 - ALMA, 6 - LSMA, 7 - HULL MA
    :param price_df: 价格数据DataFrame
    :param n: 周期参数
    :param calculation_tool: 计算工具选择序号
    :return: 平均日波动范围指标
    '''
    # 定义计算工具映射：键为序号，值为(MA计算函数, 指标名称前缀)
    tool_map = {
        0: (sma, 'SMA'),
        1: (ema, 'EMA'),
        2: (wma, 'WMA'),
        3: (ewma, 'EWMA'),
        4: (vwma, 'VWMA'),
        5: (alma, 'ALMA'),
        6: (lsma, 'LSMA'),
        7: (hull_ma, 'HULL MA')
    }
    
    # 获取对应的计算函数和名称前缀，不存在则抛出异常
    try:
        ma_func, prefix = tool_map[calculation_tool]
    except KeyError:
        raise ValueError(f"不支持的计算工具序号：{calculation_tool}")
    
    # 统一计算高低价的MA值
    high_ma = ma_func(price_df=price_df, input_mode=3, n=n, from_price=True)
    low_ma = ma_func(price_df=price_df, input_mode=4, n=n, from_price=True)
    
    # 计算平均日波动范围并重命名列
    adr = high_ma[f'{prefix}_{n}'] - low_ma[f'{prefix}_{n}']
    adr = adr.rename(columns={0: f'Average Day Range {prefix}{n}'})
    return adr

关键优化点说明 #

• 消除重复代码：原来每个if-elif分支里的调用逻辑、差值计算、列重命名都是重复的，现在通过统一逻辑处理所有情况
• 扩展更便捷：新增计算工具时，只需要在tool_map里添加对应的键值对即可，不需要新增分支
• 修复原代码的小问题：原代码里存在多处名称不匹配的错误（比如EMA分支里重命名用了SMA前缀、WMA分支里变量名写错），重构后通过统一前缀避免这类问题
• 错误处理更规范：通过捕获KeyError明确提示不支持的计算工具序号，比原代码的etc(...)更清晰

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