Update (2023-03-03)
Class decorator solution
A convenient way to solve this is by creating a reusable decorator that adds both a __get_validators__ method and a __modify_schema__ method to any given Enum class. Both of these methods are documented here.
We can define a custom validator function that will be called for our decorated Enum classes, which will enforce that only names will be turned into members and actual members will pass validation.
The schema modifier will ensure that the JSON schema only shows the names as enum options.
from collections.abc import Callable, Iterator
from enum import EnumMeta
from typing import Any, Optional, TypeVar, cast
from pydantic.fields import ModelField
E = TypeVar("E", bound=EnumMeta)
def __modify_enum_schema__(
field_schema: dict[str, Any],
field: Optional[ModelField],
) -> None:
if field is None:
return
field_schema["enum"] = list(cast(EnumMeta, field.type_).__members__.keys())
def __enum_name_validator__(v: Any, field: ModelField) -> Any:
assert isinstance(field.type_, EnumMeta)
if isinstance(v, field.type_):
return v # value is already an enum member
try:
return field.type_[v] # get enum member by name
except KeyError:
raise ValueError(f"Invalid {field.type_.__name__} `{v}`")
def __get_enum_validators__() -> Iterator[Callable[..., Any]]:
yield __enum_name_validator__
def validate_by_name(cls: E) -> E:
setattr(cls, "__modify_schema__", __modify_enum_schema__)
setattr(cls, "__get_validators__", __get_enum_validators__)
return cls
Usage
from enum import Enum
from random import choices, random
from string import ascii_lowercase
from pydantic import BaseModel
# ... import validate_by_name
# Randomly generate an enum of floats:
_members = {
name: round(random(), 1)
for name in choices(ascii_lowercase, k=3)
}
Factor = Enum("Factor", _members) # type: ignore[misc]
validate_by_name(Factor)
first_member = next(iter(Factor))
print("`Factor` members:", Factor.__members__)
print("First `Factor` member:", first_member)
class Foo(Enum):
member_a = "a"
member_b = "b"
@validate_by_name
class Bar(int, Enum):
x = 1
y = 2
class Model(BaseModel):
factor: Factor
foo: Foo
bar: Bar
class Config:
json_encoders = {Factor: lambda field: field.name}
obj = Model.parse_obj({
"factor": first_member.name,
"foo": "a",
"bar": "x",
})
print(obj.json(indent=4))
print(Model.schema_json(indent=4))
Example output:
`Factor` members: {'r': <Factor.r: 0.1>, 'j': <Factor.j: 0.9>, 'z': <Factor.z: 0.6>}
First `Factor` member: Factor.r
{
"factor": "r",
"foo": "a",
"bar": 1
}
{
"title": "Model",
"type": "object",
"properties": {
"factor": {
"$ref": "#/definitions/Factor"
},
"foo": {
"$ref": "#/definitions/Foo"
},
"bar": {
"$ref": "#/definitions/Bar"
}
},
"required": [
"factor",
"foo",
"bar"
],
"definitions": {
"Factor": {
"title": "Factor",
"description": "An enumeration.",
"enum": [
"r",
"j",
"z"
]
},
"Foo": {
"title": "Foo",
"description": "An enumeration.",
"enum": [
"a",
"b"
]
},
"Bar": {
"title": "Bar",
"description": "An enumeration.",
"enum": [
"x",
"y"
],
"type": "integer"
}
}
}
This just demonstrates a few variations for this approach. As you can see, the Factor and Bar enums are both validated by name, whereas Foo is validated by value (as a regular Enum).
Since we defined a custom JSON Encoder for Factor, the factor value is exported/encoded as the name string, while both Foo and Bar are exported by value (as a regular Enum).
Both Factor and Bar display the enum names in their JSON schema, while Foo shows the enum values.
Note that the "type": "integer" for the JSON Schema of Bar is only present because I specified int as a explicit base class of Bar and disappears, if we remove that. To further ensure consistency, we could of course also simply add "type": "string" inside our __modify_enum_schema__ function.
The only thing that is seemingly impossible right now is to also somehow register our custom way of encoding those enums inside our decorator, so that we do not need to set it in the Config or pass the encoder argument to json explicitly. That may be possible with a few changes to the BaseModel logic, but I think this would be overkill.
Original answer
Validating Enum by name
The parsing part of your problem can be solved fairly easily with a custom validator.
Since a validator method can take the ModelField as an argument and that has the type_ attribute pointing to the type of the field, we can use that to try to coerce any value to a member of the corresponding Enum.
We can actually write a more or less generalized implementation that applies to any arbitrary Enum subtype fields. If we use the "*" argument for the validator, it will apply to all fields, but we also need to set pre=True to perform our checks before the default validators kick in:
from enum import Enum
from typing import Any
from pydantic import BaseModel, validator
from pydantic.fields import ModelField
class CustomBaseModel(BaseModel):
@validator("*", pre=True)
def coerce_to_enum_member(cls, v: Any, field: ModelField) -> Any:
"""For any `Enum` typed field, attempt to """
type_ = field.type_
if not (isinstance(type_, type) and issubclass(type_, Enum)):
return v # field is not an enum type
if isinstance(v, type_):
return v # value is already an enum member
try:
return type_(v) # get enum member by value
except ValueError:
try:
return type_[v] # get enum member by name
except KeyError:
raise ValueError(f"Invalid {type_.__name__} `{v}`")
That validator is agnostic of the specific Enum subtype and it should work for all of them because it uses the common EnumType API, such as EnumType.__getitem__ to get the member by name.
The nice thing about this approach is that while valid Enum names will be turned into the correct Enum members, passing a valid Enum value still works as it did before. As does passing the member directly.
Enum names in the JSON Schema
This is a bit more hacky, but not too bad.
Pydantic actually allows us to easily customize schema generation for specific fields. This is done by adding the __modify_schema__ classmethod to the type in question.
For Enum this turns out to be tricky, especially since you want to it to be created dynamically (via the Functional API). We cannot simply subclass Enum and add our modifier method there due to some magic around the EnumType. What we can do is simply monkey-patch it into Enum (or alternatively do that to our specific Enum subclasses).
Either way, this method again gives us all we need to replace the default "enum" schema section with an array of names instead of values:
from enum import Enum
from typing import Any, Optional
from pydantic.fields import ModelField
def __modify_enum_schema__(
field_schema: dict[str, Any],
field: Optional[ModelField],
) -> None:
if field is None:
return
enum_cls = field.type_
assert isinstance(enum_cls, type) and issubclass(enum_cls, Enum)
field_schema["enum"] = list(enum_cls.__members__.keys())
# Monkey-patch `Enum` to customize schema modification:
Enum.__modify_schema__ = __modify_enum_schema__ # type: ignore[attr-defined]
And that is all we need. (Mypy will complain about the monkey-patching of course.)
Full demo
from enum import Enum
from random import choices, random
from string import ascii_lowercase
from typing import Any, Optional
from pydantic import BaseModel, validator
from pydantic.fields import ModelField
def __modify_enum_schema__(
field_schema: dict[str, Any],
field: Optional[ModelField],
) -> None:
if field is None:
return
enum_cls = field.type_
assert isinstance(enum_cls, type) and issubclass(enum_cls, Enum)
field_schema["enum"] = list(enum_cls.__members__.keys())
# Monkey-patch `Enum` to customize schema modification:
Enum.__modify_schema__ = __modify_enum_schema__ # type: ignore[attr-defined]
class CustomBaseModel(BaseModel):
@validator("*", pre=True)
def coerce_to_enum_member(cls, v: Any, field: ModelField) -> Any:
"""For any `Enum` typed field, attempt to """
type_ = field.type_
if not (isinstance(type_, type) and issubclass(type_, Enum)):
return v # field is not an enum type
if isinstance(v, type_):
return v # value is already an enum member
try:
return type_(v) # get enum member by value
except ValueError:
try:
return type_[v] # get enum member by name
except KeyError:
raise ValueError(f"Invalid {type_.__name__} `{v}`")
# Randomly generate an enum of floats:
_members = {
name: round(random(), 1)
for name in choices(ascii_lowercase, k=3)
}
Factor = Enum("Factor", _members) # type: ignore[misc]
first_member_name = next(iter(Factor)).name
print("Random `Factor` members:", Factor.__members__)
print("First member:", first_member_name)
class Model(CustomBaseModel):
factor: Factor
foo: str
bar: int
class Config:
json_encoders = {Factor: lambda field: field.name}
obj = Model.parse_obj({
"factor": first_member_name,
"foo": "spam",
"bar": -1,
})
print(obj.json(indent=4))
print(Model.schema_json(indent=4))
Output:
Random `Factor` members: {'a': <Factor.a: 0.9>, 'q': <Factor.q: 0.6>, 'e': <Factor.e: 0.8>}
First member: a
{
"factor": "a",
"foo": "spam",
"bar": -1
}
{
"title": "Model",
"type": "object",
"properties": {
"factor": {
"$ref": "#/definitions/Factor"
},
"foo": {
"title": "Foo",
"type": "string"
},
"bar": {
"title": "Bar",
"type": "integer"
}
},
"required": [
"factor",
"foo",
"bar"
],
"definitions": {
"Factor": {
"title": "Factor",
"description": "An enumeration.",
"enum": [
"a",
"q",
"e"
]
}
}
}
Notes
I chose this super weird way of randomly generating an Enum just for illustrative purposes. I wanted to show that both validation and schema generation still work fine in that case. But in practice I would assume that the names actually don't change that drastically every time the program is run. (At least I hope they don't for the sake of your users.)
The value of factor is still a regular Enum member, so obj.factor.value will still give us 0.9 (for this random example).
The validator will obviously prevent invalid names/values to be passed. You can make it more specific, if you like or restrict it to only deal with str arguments assuming them to be Enum member names and delegate the rest to Pydantic's default validator. As it is written right now, it essentially replaces that default Enum validator.
Any other schema modifications (such as the description) can be done according to the docs I linked as well.
