Python如何检查pandas DataFrame列是否存在（按列名/索引号）及代码优化

原代码可优化点 #

• 函数命名不合理：read过于通用，和pandas内置read_*系列函数易混淆，不符合你选择标签列的使用场景，表意不清晰
• 存在冗余死代码：函数末尾的return data[column]永远不会被执行，且data变量未定义，属于无效代码
• 整数输入判定逻辑鲁棒性差：仅通过lstrip('-').isdecimal()判断无法处理带前后空格的数字输入、误输入小数等场景，容易触发报错或者逻辑错误
• 错误提示不够友好：没有给用户明确的有效输入范围参考，用户不知道可选列名、合法索引区间，试错成本高
• 重复计算：多次执行int(column)、df.shape[1]调用，可提前缓存结果减少重复计算
• 取值逻辑冗余：不需要分开调用iloc和loc，pandas原生支持直接用列名/整数下标取列，写法可以更简洁

优化后代码 #

import pandas as pd

def select_label_column(df):
    # 提前缓存列数、列名列表，减少重复计算
    col_count = df.shape[1]
    col_names = df.columns.tolist()
    while True:
        input_val = input("请输入要选择的列名或列索引号：").strip()
        # 先尝试转整数处理索引输入
        try:
            col_idx = int(input_val)
            # 判断索引是否合法
            if -col_count <= col_idx < col_count:
                return df.iloc[:, col_idx]
            else:
                print(f"列索引不存在，有效索引范围为 [{-col_count}, {col_count-1}]，请重试\n")
                continue
        except ValueError:
            # 转整数失败，判定为列名输入
            if input_val in col_names:
                return df[input_val]
            else:
                print(f"列名不存在，当前可用列名为：{col_names}，请重试\n")
                continue

# 测试用例
if __name__ == "__main__":
    df = pd.DataFrame({'A': [1,2,3], 'B':['a', 'b', 'c']})
    selected_col = select_label_column(df)
    print("选中的列内容：\n", selected_col)

优化效果说明 #

• 支持处理带前后空格的输入，用户不小心打了空格也能正常识别
• 错误提示给出明确的参考范围，降低用户输入难度
• 用异常捕获处理整数转换逻辑，覆盖更多异常输入场景，不会因为奇怪的输入直接崩溃
• 函数命名贴合你计算香农熵选标签列的场景，代码可读性更高
• 删除了无效死代码，避免潜在变量未定义报错
• 新增测试代码判定逻辑，导入该函数到其他项目时不会自动执行测试逻辑

内容的提问来源于stack exchange，提问作者Александр Ермак