自定义函数开发：如何强制传入指定列结构的pd.DataFrame？

你可以在函数开头加入列名验证逻辑，通过对比列名集合确保传入的DataFrame恰好包含且仅包含Col_A和Col_B两列，同时给出清晰的错误提示。修改后的函数如下：

import pandas as pd

def MyFn(DF: pd.DataFrame) -> float:
    # 定义必须存在的列名集合
    required_columns = {'Col_A', 'Col_B'}
    # 获取传入DataFrame的列名集合
    df_columns = set(DF.columns)
    
    # 检查列名是否完全匹配
    if df_columns != required_columns:
        # 计算缺失和多余的列
        missing_cols = required_columns - df_columns
        extra_cols = df_columns - required_columns
        
        # 构建错误提示信息
        error_message = []
        if missing_cols:
            error_message.append(f"缺少必要列: {', '.join(missing_cols)}")
        if extra_cols:
            error_message.append(f"存在多余列: {', '.join(extra_cols)}")
        
        raise ValueError('; '.join(error_message))
    
    # 原业务逻辑
    return DF['Col_A'].values[1] - DF['Col_B'].values[1]

验证逻辑说明 #

1. 集合对比：用集合存储列名，无需关心列的传入顺序，只要列名完全一致即通过验证。
2. 精准错误提示：分别识别缺失的必要列和多余的无效列，让用户能快速修正问题。
3. 标准异常抛出：使用ValueError抛出参数验证失败的异常，符合Python参数校验的惯例。

测试示例 #

# 符合要求的输入
valid_df = pd.DataFrame({'Col_A': [10, 25, 30], 'Col_B': [5, 10, 15]})
print(MyFn(valid_df))  # 输出 15

# 缺少列的情况（会报错）
invalid_df1 = pd.DataFrame({'Col_A': [1, 2, 3]})
MyFn(invalid_df1)  # 抛出 ValueError: 缺少必要列: Col_B

# 存在多余列的情况（会报错）
invalid_df2 = pd.DataFrame({
    'Col_A': [1, 2, 3],
    'Col_B': [4, 5, 6],
    'Col_C': [7, 8, 9]
})
MyFn(invalid_df2)  # 抛出 ValueError: 存在多余列: Col_C

内容的提问来源于stack exchange，提问作者Bogaso