Как правильно протестировать dtypes Pandas в рамках датафреймов?

Question

Цель: создать функцию, которая может сопоставить данные dtypes с заранее заданным сценарием типа данных.

Описание: Я хочу иметь возможность классифицировать данные наборы данных на основе их отнесения к предопределенным типам сценариев.

Ниже приведены два примера наборов данных (df_a и df_b). df_a имеет только dtypes, равные 'object', тогда как df_b имеет и 'object', и 'int64':

# scenario_a
data_a = [['tom', 'blue'], ['nick', 'green'], ['julia', 'red']]  
df_a = pd.DataFrame(data, 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, columns = ['Name', 'Age'])

Я хочу, чтобы можно было автоматически определить, какой сценарий основан на функции:

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

Это то, что я имею до сих пор, но не получаю результатов, на которые я надеялся.

При использовании функции scenario(df_a) я ищу результат, равный scenario_a, а когда я передаю df_b, я ищу функцию, которая сможет правильно определить, какой сценарий должен быть.

Любая помощь будет оценена.

Answer 1

Вот один из подходов. Создайте dict scenarios, с ключами, отсортированными tuple с предопределенными dtypes, а значением будет то, что вы хотели бы получить функцией.

Используя ваш пример, что-то вроде:

# 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_a = tuple(sorted(df_a.dtypes.unique()))
scenario_b = tuple(sorted(df_b.dtypes.unique()))

scenarios = {
    scenario_a: 'scenario_a',
    scenario_b: 'scenario_b'
}

print(scenarios)

# scenarios:
# {(dtype('O'),): 'scenario_a', (dtype('int64'), dtype('O')): 'scenario_b'}

def scenario(data):
    dtypes = tuple(sorted(data.dtypes.unique()))
    return scenarios.get(dtypes, None)

scenario(df_a)
# 'scenario_a'

scenario(df_b)
# scenario_b