如何实现Pandas DataFrame数据类型场景匹配函数？求修正方案

我需要编写一个函数，根据DataFrame的列数据类型，自动将其归类到预定义的场景中。具体要求：

• scenario_a：DataFrame的所有列均为object类型（比如示例中的df_a）
• scenario_b：DataFrame同时包含object和int64类型（比如示例中的df_b）

我的示例数据集如下：

import pandas as pd

# scenario_a 对应数据集
data_a = [['tom', 'blue'], ['nick', 'green'], ['julia', 'red']]
df_a = pd.DataFrame(data_a, columns=['Name','Color'])
df_a['Color'] = df_a['Color'].astype('object')

# scenario_b 对应数据集
data_b = [['tom', 10], ['nick', 15], ['julia', 14]]
df_b = pd.DataFrame(data_b, columns=['Name', 'Age'])

我尝试编写了以下函数，但无法得到预期结果：调用scenario(df_a)应返回scenario_a，调用scenario(df_b)应返回scenario_b，目前函数运行存在问题，请求帮助。

我的错误代码：

import pandas as pd
import numpy as np
def scenario(data):
    if data.dtypes.str.contains('object'):
        return scenario_a
    if data.dtypes.str.contatin('object', 'int64'):
        return scenario_b

问题分析 #

你的函数存在三个核心问题：

1. 拼写错误：第二个判断中的contatin是笔误，正确写法应为contains
2. 布尔判断逻辑错误：data.dtypes.str.contains('object')返回的是一个布尔Series（每列对应一个True/False），直接用if判断会触发ValueError——Python无法将Series直接当作单个布尔值处理
3. 多类型判断方式错误：contains方法不支持同时传入多个类型字符串，你需要通过列类型的集合来判断组合关系

修正后的解决方案 #

import pandas as pd

# 定义场景标识（可根据需求替换为其他类型，比如函数、类等）
SCENARIO_A = "scenario_a"
SCENARIO_B = "scenario_b"

def determine_scenario(df):
    # 将所有列的类型转为字符串，存入集合方便判断组合
    dtype_set = set(df.dtypes.astype(str))
    
    # 判断scenario_a：仅包含object类型
    if dtype_set == {'object'}:
        return SCENARIO_A
    # 判断scenario_b：同时包含object和int64类型，无其他类型
    elif dtype_set == {'object', 'int64'}:
        return SCENARIO_B
    # 处理未匹配到的场景
    else:
        return "unknown_scenario"

测试验证 #

# 测试df_a
print(determine_scenario(df_a))  # 输出: scenario_a
# 测试df_b
print(determine_scenario(df_b))  # 输出: scenario_b

灵活调整逻辑 #

如果你的场景允许存在其他类型（比如scenario_b只要同时有object和int64即可，不管是否有其他类型），可以修改判断条件为：

elif 'object' in dtype_set and 'int64' in dtype_set:
    return SCENARIO_B

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