I have written a function which asks a user for column name (ex. 'Age') or column number (0, 1, ... or -1, -2, ...) and returns it if exists. I would like to know if my solution can be improved in terms of code design.
To clarify, I need this piece of code for another function which calculates Shannon Entropy on the dataframes for which label-column should be chosen manually.
import pandas as pd
df = pd.DataFrame({'A': [1,2,3], 'B':['a', 'b', 'c']})
def read(df):
while True:
column = input("Please, enter column name or number:")
if column.lstrip('-').isdecimal():
if (-df.shape[1] > int(column)) or (int(column) >= df.shape[1]):
print('Such column does not exist. Please, try again. \n')
continue
else:
return df.iloc[:, int(column)]
elif column not in df.columns:
print('Such column does not exist. Please, try again. \n')
continue
else:
return df.loc[:, column]
return data[column]
read(df)
Aucun commentaire:
Enregistrer un commentaire