How to make the columns selections in dataframe when passed few columns as list

I have the function contains 3 parameters:

def foo(df, columns, additional_col=None):
       df = df[columns + additoinal columns]

if additional_col parameter then only it should append it to columns else it should keep columns as column selection

Example:

columns = ["A", "B", "C", "D"]
addtional_col = ["X", "Y"]

if additional_col is passed while calling the function foo then column selection would be df["A", "B", "C", "D", "X", "Y"] elseif additional_col is None then df["A", "B", "C", "D"]

tried join, map and split but couldn`t achieve the desire output. Need help on immediate basis. Thanks

Solution

Firstly, you will need to make sure that you make a copy of the columns list to prevent unexpected side effects of extending the original list.
If additional_col has items in the list it will equate to True when used in an if-statement.
So if additional_col has items, you can extend the columns list using the extend() function.
If it does not have items, then just use the original columns list.

Here is the code:

Code:

def foo(df, columns, additional_col=None):
    columns = list(columns)
    if additional_col:
        columns.extend(additional_col)
        df = df[columns]
    else:
        df = df[columns]
    
    return df
        

data = pd.DataFrame({"A":[1,2,3], "B":[4,5,6], "C":[7,8,9], "X":['a','b','c'], "Y":['d','e','f']})
    
cols = ["A","B","C"]
a = ["X","Y"]

print(foo(data, cols,a))

print("-------------------")

print(foo(data, cols))

Output:

   A  B  C  X  Y
0  1  4  7  a  d
1  2  5  8  b  e
2  3  6  9  c  f
-------------------
   A  B  C
0  1  4  7
1  2  5  8
2  3  6  9