Получение аргумента typing.Optional в Python

На чтение 3 мин Опубликовано 17.09.2024

Я хотел бы создать типизированный DataFrame из класса Pydantic BaseModel, назовем его MyModel, который имеет необязательные поля. Поскольку я создаю несколько экземпляров MyModel, некоторые из них будут иметь необязательные поля со значениями None, и если я инициализирую DataFrame с такими строками, они могут иметь непоследовательные типы колонок. Я хотел бы привести Optional[TypeX] к TypeX, например:

import pydantic
import pandas as pd
import numpy as np
from typing import Optional

class MyModel(pydantic.BaseModel):
   thisfield: int
   thatfield: Optional[str]
   ...

col_types = {kk: ff.annotation for kk, ff in MyModel.model_fields.items()}

pd.DataFrame(np.empty(0, dtype=[tuple(tt) for tt in col_types.items()]))

Это вызывает ошибку TypeError: Cannot interpret 'typing.Optional[str]' as a data type.

Мне нужна функция или метод для Optional[X] -> X. Есть ли другие предложения, кроме использования repr с регулярными выражениями?

Ответ

Для решения задачи, вы можете создать вспомогательную функцию, которая будет преобразовывать аннотации типов с использованием Optional[X] в простые типы X. В вашей ситуации мы используем библиотеку Pydantic для определения модели данных и библиотеку pandas для создания DataFrame. Проблема заключается в том, что pandas не может правильно интерпретировать типы, содержащие Optional.

Ниже приведен полный код, который решает вашу проблему, включая создание DataFrame с правильными типами данных:

import pydantic
import pandas as pd
import numpy as np
from typing import Optional, Any

class MyModel(pydantic.BaseModel):
    thisfield: int
    thatfield: Optional[str]

def strip_optional(annotation: Any) -> Any:
    """Функция для удаления Optional из аннотации типа."""
    if getattr(annotation, '__origin__', None) is Optional:
        return annotation.__args__[0]  # Вернуть первый аргумент, если это Optional
    return annotation

# Получаем типы столбцов из модели MyModel
col_types = {kk: strip_optional(ff.annotation) for kk, ff in MyModel.model_fields.items()}

# Создаем DataFrame с использованием полученных типов
dtype = {kk: np.dtype(ff) if ff in (int, float) else 'object' for kk, ff in col_types.items()}
data = np.empty(0, dtype=[(name, dtype[name]) for name in dtype.keys()])

# Инициализируем DataFrame
df = pd.DataFrame(data)

# Проверим получившийся DataFrame
print(df.dtypes)

Объяснение

Модель Pydantic: Мы определяем класс MyModel с обязательным целочисленным полем и необязательным строковым полем.
Функция strip_optional: Эта функция принимает аннотацию типа и проверяет, является ли она Optional. Если да, то она возвращает внутренний тип. Если нет, возвращает саму аннотацию. Мы используем getattr и __origin__, чтобы определить, является ли аннотация Optional.
Получение типов столбцов: Мы создаем словарь col_types, который преобразует аннотации типов полей из модели MyModel, использовав функцию strip_optional.
Создание DataFrame: Мы создаем массив с пустыми данными и используем его для инициализации DataFrame, правильно указывая типы данных.
Проверка: В конце кода мы выводим типы данных в DataFrame, чтобы убедиться, что они установлены корректно.

Таким образом, вы сможете корректно обрабатывать Optional типы в ваших моделях и инициализировать DataFrame со стандартными типами данных.