def add_nim_input_columns_from_tuple_or_dict()

in spotify_confidence/analysis/frequentist/nims_and_mdes.py [0:0]

• 21 lines of code
• 9 McCabe index (conditional complexity)

def add_nim_input_columns_from_tuple_or_dict(df, nims: NIM_TYPE, mde_column: str) -> DataFrame:
    if type(nims) is tuple:
        return df.assign(**{NIM_COLUMN_DEFAULT: nims[0]}).assign(**{PREFERRED_DIRECTION_COLUMN_DEFAULT: nims[1]})
    elif type(nims) is dict:
        nim_values = {key: value[0] for key, value in nims.items()}
        nim_preferences = {key: value[1] for key, value in nims.items()}
        return df.assign(**{NIM_COLUMN_DEFAULT: lambda df: df.index.to_series().map(nim_values)}).assign(
            **{PREFERRED_DIRECTION_COLUMN_DEFAULT: lambda df: df.index.to_series().map(nim_preferences)}
        )
    elif nims is None or not nims:
        return df.assign(**{NIM_COLUMN_DEFAULT: None}).assign(
            **{
                PREFERRED_DIRECTION_COLUMN_DEFAULT: (
                    None
                    if PREFERRED_DIRECTION_COLUMN_DEFAULT not in df or mde_column is None
                    else df[PREFERRED_DIRECTION_COLUMN_DEFAULT]
                )
            }
        )
    else:
        return df