Как использовать объекты pycountry.db.Country в качестве индекса pd.DataFrame?

Объяснение

Pandas рассматривает ваш объект как объект, похожий на список, поэтому вы не можете использовать его в качестве ключа для loc, поскольку он попытается перебрать объекты в списке.

>>> from pandas.core.dtypes.common import is_list_like, is_scalar
>>> is_scalar(belgium)
False
>>> is_list_like(belgium)
True

См. Какой тип данных в Python считается списком? подробнее о is_list_like

Обходной путь

Интересно, это работает:

>>> df.loc[[belgium]].iloc[0]
population    11429336
Name: Country(alpha_2='BE', alpha_3='BEL', flag='🇧🇪', name='Belgium', numeric='056', official_name='Kingdom of Belgium'), dtype: int64

поэтому, если вы действительно хотите использовать объект в качестве индекса, вы можете обойти это с помощью этого.

Или, что еще более смешно, сделать объект неперебираемым:

>>> belgium.__iter__ = None
>>> df.loc[belgium]
population    11429336
Name: Country(__iter__=None, alpha_2='BE', alpha_3='BEL', flag='🇧🇪', name='Belgium', numeric='056', official_name='Kingdom of Belgium'), dtype: int64

Но я уверен, что это нарушит некоторые другие функции вашего кода, поскольку __iter__, похоже, реализовано в Country, чтобы его можно было легко привести к dict.

Рекомендация

Использование объектов в качестве индекса, возможно, не является лучшим решением, если ваш набор данных большой. Я бы порекомендовал вместо этого использовать, например, alpha_3 в качестве индекса и хранить объект в отдельном столбце. Вы по-прежнему сможете избежать двусмысленности, но не столкнетесь с проблемами со слишком сложными типами индексов.

data: list[dict] = [
    {"index": aruba.alpha_3, "country": aruba, "population": 106_203},
    {"index": belgium.alpha_3, "country": belgium, "population": 11_429_336},
    {"index": canada.alpha_3, "country": canada, "population": 37_058_856},
]

df: pd.DataFrame = pd.DataFrame(data)
df.set_index("index", inplace=True)

assert df.loc[belgium.alpha_3]["country"] == belgium

Хотя я согласен с рекомендацией @Teemu Risikko не использовать объекты в качестве индекса, его превосходный ответ также заставил меня понять, почему ваш код работает в версиях pycountry <= 23.12.7.

В версии 23.12.7 ваши объекты будут pycountry.db.Data, там все еще нет метода __iter__ (см. 23.12.7 db.py L8).

С 23.12.11 объекты pycountry.db.Country наследуются от Data с __iter__ (см. 23.12.11 db.py L32 , L38).

Например, если вы хотите настоять:

• перейти на версию 23.12.7 (используя pip install pycountry==23.12.7 или эквивалент) и;
• обновите подсказки типов (pycountry.db.Country -> pycountry.db.Data)

import pandas as pd
import pycountry

belgium: pycountry.db.Data = pycountry.countries.get(alpha_3 = "BEL")

data: list[dict] = [
    {"country": belgium, "population": 11_429_336}
]

df: pd.DataFrame = pd.DataFrame(data)
df.set_index("country", inplace=True)

df.loc[belgium]

Выход:

population    11429336
Name: Country(alpha_2='BE', alpha_3='BEL', flag='🇧🇪', name='Belgium', numeric='056', official_name='Kingdom of Belgium'), dtype: int64

print(pycountry.__version__)
# 23.12.7

Благодаря ответу Теему Рисикко , а также ответу Ouroboros1, я создал простой Country(StrEnum) класс, который позволяет мне комбинировать:

• проверка (можно создать экземпляр только действительных pycountry стран),
• объединение триграммы страны (просто) и названия (полезно для отображения),
• выбор из объекта Country (любая строка недопустима).

Допустимые экземпляры Enum вычисляются автоматически.

import pandas as pd
import pycountry
from enum import StrEnum

Country = StrEnum(
    value = "Country",
    names = {country.alpha_3: country.name for country in pycountry.countries},
    module=__name__,
)

aruba = Country.ABW  # <-- One can select a country by its alpha-3 code
belgium = Country.BEL
canada = Country("Canada")  # <-- One can select a country by (standardized) name too

data: list[dict] = [
    {"country": aruba, "population": 106_203},
    {"country": belgium, "population": 11_429_336},
    {"country": canada, "population": 37_058_856},
]

df: pd.DataFrame = pd.DataFrame(data)
df.set_index("country", inplace=True)

assert df.index[1] == belgium
assert df.index[1] is belgium

belgium_data = df.loc[belgium]