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
CodePudding user response:
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 toTrue
when used in anif-statement
.So if
additional_col
has items, you can extend thecolumns
list using theextend()
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