CODEWITHSUNDEEP
pythongenericsenumspydantic
TomTom101
TomTom101

Reputation: 6892

Include possible enum values in pydantic model

I have a model ModelWithEnum that holds an enum value. In addition to that value, I want the model to output all possible values from that enum (those enums are range-like, e.g. A,B,C; A,B,C,D etc.).

I made EnumModel a GenericModel since I don't want a new model for every Enum class. I am doing my first steps with generics and not sure, whether I am on the right track.

Here's what I want to achieve

# `ModelWithEnum` holds a `enum` field of type `EnumsAC`,
# so possible enum values would be "A", "B", "C"
model = ModelWithEnum(enum={"value": "A"})

assert model.enum.value == "A"
assert model.enum.possible_values == ["A", "B", "C"]

Full sample code:

from enum import Enum
from typing import Generic, List, Optional, TypeVar
from pydantic.generics import GenericModel
from pydantic import BaseModel

EnumT = TypeVar("EnumT")


class EnumsAC(str, Enum):
    A = "A"
    B = "B"
    C = "C"


class EnumsAE(str, Enum):
    """Another Enums class"""

    A = "A"
    B = "B"
    C = "C"
    D = "D"
    E = "E"


class EnumModel(GenericModel, Generic[EnumT]):
    value: EnumT
    # How can I auto-fill this list with all possible values from EnumT
    # like if I hardcoded e.g. `list(EnumsAC)`
    possible_values: List[str] = []


class ModelWithEnum(BaseModel):
    enum: EnumModel[EnumsAC]

# `ModelWithEnum` holds a `enum` field of type `EnumsAC`,
# so possible enum values would be "A", "B", "C"
model = ModelWithEnum(enum={"value": "A"})

assert model.enum.value == "A" # ok
assert model.enum.possible_values == ["A", "B", "C"] # fails

Upvotes: 1

Views: 6498

Answers (1)

alex_noname
alex_noname

Reputation: 32233

You could use root_validator for autocomplete possible values. Using your model as an example:

class EnumModel(GenericModel, Generic[EnumT]):
    value: EnumT
    possible_values: List[str] = []

    class Config:
        validate_assignment = True

    @root_validator
    def root_validate(cls, values):
        values["possible_values"] = [item for item in values['value'].__class__]
        return values

Upvotes: 1

Related Questions