pandas adding an optional suffix to a column name as a parameter in a function

Question

Objective: I have a pandas dataframe with many column names that are identical except for a suffix. I would like to be able to reference the column in a function by passing the suffix via a parameter. Example:

import pandas as pd
df_test = pd.DataFrame({'HHLD_ID':[1,2,3,4,5],
                       'sales':[25,50,25,25,50],
                       'sales_sk':[1,2,1,1,2],
                       'sales_rf':[1,2,1,1,5]
                       })

def total_print(df,label,suffix=''):  
    print(f'{label} Total dollars, {df[sales+suffix].sum()}')

total_print(df_test,"Your rf sales are :","_rf")
total_print(df_test,"Your sales are :")

Expected Answer:
Your rf sales are: 10
Your sales are: 175

Answer 1

IIUC, you could use string format:

import pandas as pd

df_test = pd.DataFrame({'HHLD_ID': [1, 2, 3, 4, 5],
                        'sales': [25, 50, 25, 25, 50],
                        'sales_sk': [1, 2, 1, 1, 2],
                        'sales_rf': [1, 2, 1, 1, 5]
                        })

sales = "sales{}".format


def total_print(df, label, suffix=''):
    print(f'{label} Total dollars, {df[sales(suffix)].sum()}')


total_print(df_test, "Your rf sales are :", "_rf")
total_print(df_test, "Your sales are :")

Output

Your rf sales are : Total dollars, 10
Your sales are : Total dollars, 175

As an alternative you could also use f-strings directly:

def total_print(df, label, suffix=''):
    print(f'{label} Total dollars, {df[f"sales{suffix}"].sum()}')