[FIXED] Python How to check presence of a column in dataframe by it's name or number

Issue

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)

Solution

The columns are available in df.columns which can be used to get the data you want. If the column isn't in df.columns, try converting it to an int to index df.columns and use an exception handler to deal with the misses.

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 not in df.columns:
            try:
                column = df.columns[int(column)]
            except (IndexError, ValueError):
                print(f"Column {column!r} does not exist, Please try again.")
                continue
        break
    return df.loc[:, column]

print(read(df))

Answered By - tdelaney