What to use instead of if-elif inside a function to choose between multiple functions?-CodePudding

I am trying to create my own module that will calculate indicators, and I am implementing a function like this:

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:
    :param n:
    :param calculation_tool:
    :return:
    '''
    if calculation_tool == 0:
        sma_high = sma(price_df=price_df, input_mode=3, n=n, from_price=True)
        sma_low = sma(price_df=price_df, input_mode=4, n=n, from_price=True)
        adr = sma_high[f'SMA_{n}'] - sma_low[f'SMA_{n}']
        adr.rename(columns={0: f'Average Day Range SMA{n}'})
        return adr
    elif calculation_tool == 1:
        ema_high = ema(price_df=price_df, input_mode=3, n=n, from_price=True)
        ema_low = ema(price_df=price_df, input_mode=4, n=n, from_price=True)
        adr = ema_high[f'EMA_{n}'] - ema_low[f'EMA_{n}']
        adr.rename(columns={0: f'Average Day Range SMA{n}'})
        return adr
    elif calculation_tool == 2:
        ema_high = wma(price_df=price_df, input_mode=3, n=n, from_price=True)
        ema_low = wma(price_df=price_df, input_mode=4, n=n, from_price=True)
        adr = ema_high[f'EMA_{n}'] - ema_low[f'EMA_{n}']
        adr.rename(columns={0: f'Average Day Range SMA{n}'})
        return adr
    elif calculation_tool == 3:
        ewma_high = ewma(price_df=price_df, input_mode=3, n=n, from_price=True)
        ewma_low = ewma(price_df=price_df, input_mode=4, n=n, from_price=True)
        adr = ewma_high[f'EMA_{n}'] - ewma_low[f'EMA_{n}']
        adr.rename(columns={0: f'Average Day Range SMA{n}'})
        return adr
    
    etc(...)

Is there a more convinient way to do the same but without if-elif hell?

CodePudding user response：

ok I have no idea what ema, sma, wma or ewma functions are but I can see that you are using calculation_tool to select suited functions for your average_day_range function

you could map your functions as below:

function_dict = {
    0: {
        "func": sma,
        "name": "SMA"
    },
    1: {
        "func": ema,
        "name": "EMA"
    },
    2: {
        "func": wma,
        "name": "WMA"
    },
    3: {
        "func": ewma,
        "name": "EWMA"
    },
}

then you can refactor your current code as this:

def average_day_range(price_df: PandasDataFrame, n: int = 14, calculation_tool: int = 0):
    function, name = function_dict[calculation_tool]["func"], function_dict[calculation_tool]["name"]
    high_var = function(price_df=price_df, input_mode=3, n=n, from_price=True)
    low_var = function(price_df=price_df, input_mode=4, n=n, from_price=True)
    adr = high_var[f'{name}_{n}'] - low_var[f'{name}_{n}']
    adr.rename(columns={0: f'Average Day Range {name}{n}'})
    return adr

magic Happens here at function, name = function_dict[calculation_tool]["func"], function_dict[calculation_tool]["name"]

you are essentially mapping target functions pointer to your function variable its like defining an alias

for better understanding keep in mind that you can do something like this:

my_print = print
my_print("Hello World")

CodePudding user response：

You can store the functions/strings that change in a giant list (mapping the calculation tool value to the functions/strings it uses). For example, when calculation_tool=0, it uses the sma() function and the format-string f'SMA_{n}'. You'll need to add new elements to the list for each calculation tool, and possible more values to the tuples within the list for other functions/strings/objects that change, but this allows you to make your code much more generic.

map_n_to_tool = [(sma, 'SMA_{n}'), ...]


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:
    :param n:
    :param calculation_tool:
    :return:
    '''
    func, col = map_n_to_tool[n]
    high = func(price_df=price_df, input_mode=3, n=n, from_price=True)
    low = func(price_df=price_df, input_mode=4, n=n, from_price=True)
    adr = sma_high[col.format(n)] - sma_low[col.format(n)]
    adr.rename(columns={0: f'Average Day Range SMA{n}'})
    return adr