Example features container using Pydantic

animal_features_example.py

import typing as T

import numpy as np
from pydantic import BaseModel
from pydantic.types import conint, confloat


class FeatureIsNoneError(Exception):
    pass


class AnimalFeatures(BaseModel):
    cat_1: T.Optional[conint(strict=True, ge=0, le=1)] = None
    cat_2: T.Optional[conint(strict=True, ge=0, le=1)] = None
    dog: T.Optional[confloat(strict=True, ge=0.5, le=11.5)] = None
    bird: T.Optional[confloat(strict=True, ge=-1.27, le=17.33)] = None

    class Config:
        validate_assignment = True

    def set_categorical_features(
        self, 
        prefix: str, 
        positive_category: T.Union[str, int], 
        sep: str = "_"
    ):
        found = False
        for field in self.__annotations__:
            if field == f"{prefix}{sep}{positive_category}":
                setattr(self, field, 1)
                found = True
                continue

            if prefix in field:
                setattr(self, field, 0)

        if not found:
            raise ValueError(
                f"Could not find the positive category {prefix}{sep}{positive_category}"
            )

    def set_bulk_features(self, mapping: T.Dict):
        for field, value in mapping.items():
            setattr(self, field, value)

    @property
    def numpy_array(self) -> np.ndarray:
        vals = []
        for field in self.__annotations__:
            attr = getattr(self, field)
            if attr is None:
                raise FeatureIsNoneError(f"{field} value cannot be None")
            vals.append(attr)
        return np.array([vals])