首页 > 解决方案 > Python如何通过名称或编号检查数据框中是否存在列

问题描述

我编写了一个函数,它询问用户列名(例如'Age')或列号(0、1、...或-1、-2、...)并在存在时返回它。我想知道我的解决方案是否可以在代码设计方面得到改进。

为了澄清,我需要这段代码用于另一个函数,该函数计算应该手动选择标签列的数据帧上的香农熵。

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)

标签: pythonpandasdesign-patterns

解决方案


可用的列df.columns可用于获取所需的数据。如果该列不在 中df.columns,请尝试将其转换int为索引df.columns并使用异常处理程序来处理未命中。

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))

推荐阅读