Create a python function that accepts dataframe and column(s)

Question

I'm trying to create a function that accepts dataframe columns.

Something along the lines of the below pseudo code..

def tb_mend_format(df, col):
   
    if df[col][:3] == 'TBK':
        return 'TB ' + df[col][7:]
    else:
        return df[col]

Is it possible to then pass a dataframe and column(s) in the below fashion as required? The plan is to create a reusable function.

tb_mend_format(df1, df['Key'])        

Thanks!

Answer 1

Probably you meant something like this?

def tb_mend_format(df, col: Union[pd.Series, str]):

    if isinstance(col, pd.Series):
         series = col
    else:
         series = df[col]

    if series[:3] == 'TBK':
        return 'TB ' + series[7:]
    else:
        return series

# both works
tb_mend_format(df, 'Key')
tb_mend_format(df, df['Key'])

This code, however, won't work if col is a list/tuple of str (column names).

In Pandas, a column is represented by a pd.Series object and I guess you meant column by the "name" of a column or Series, in str type (or list/tuple of str for multiple columns).