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

Question

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?

Object-oriented programming solves this rather elegantly, but there's a learning curve. Each `elif` would be a subclass with different parameters (and down the line even possibly different implementation under the hood). — tripleee, Dec 17 '22 at 17:58
It looks like you could store `ewma` and such in a dictionary (`{1: sma, 2: ema}`), then do a lookup of the dictionary. It would be a bit more complex than that because you also need to save a `_high` and `_low` dataframe, but that wouldn't complicate things too much. — Carcigenicate, Dec 17 '22 at 18:01
the `ewma_high = ...`, `ewma_low = ...` and `adr.rename ...` could be outside the if-elif block which would simplify this function quite a bit already, there's no reason to have all this duplicated code. And as @tripleee said, if this repetition is a pattern that shows up in other places as well it might be worth investing in the object oriented approach. — Janilson, Dec 17 '22 at 18:02
Have a look at: [Replacements for switch statement in Python?](https://stackoverflow.com/questions/60208/replacements-for-switch-statement-in-python) — Maurice Meyer, Dec 17 '22 at 18:06

Krishnan Shankar · Answer 1 · 2022-12-17T18:15:40.833

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

Thanks alot, I didn't knew how to do this :) – Jakub Szurlej Dec 17 '22 at 18:33 — Jakub Szurlej, Dec 17 '22 at 18:33

score 1 · Accepted Answer · answered Dec 17 '22 at 18:13

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

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

2 Answers2