在 pydantic 模型中初始化 Literal 枚举

Question

pydantic 很好地支持常规枚举，并且可以使用枚举 instance 和枚举 value 来初始化枚举类型字段：

from enum import Enum
from pydantic import BaseModel

class MyEnum(Enum):
    FOO = 'foo'
    BAR = 'bar'

class MyModel(BaseModel):
    x: MyEnum

MyModel(x=MyEnum.BAR)  # Enum instance, works
MyModel(x='foo')       # Enum value, works
MyModel(x='?')         # Fails, as expected

pydantic 也支持typing.Literal:

from typing import Literal
from pydantic import BaseModel

class MyModel(BaseModel):
    x: Literal['foo']

MyModel(x='foo')  # Works
MyModel(x='bar')  # Fails, as expected

现在我想结合枚举和文字，即强制字段值等于一个特定的枚举实例。但是，我仍然希望能够传递正确的枚举 value（即不仅是正确的枚举 instance），这似乎不起作用：

from enum import Enum
from typing import Literal
from pydantic import BaseModel

class MyEnum(Enum):
    FOO = 'foo'
    BAR = 'bar'

class MyModel(BaseModel):
    x: Literal[MyEnum.FOO]

MyModel(x=MyEnum.FOO)  # Enum instance, works
MyModel(x=MyEnum.BAR)  # Fails, as expected
MyModel(x='foo')       # Enum value, fails but I'd like it to work

Answer 1

尝试使用 基于字符串的 Enum。像这样：

from enum import Enum
from typing import Literal
from pydantic import BaseModel


class MyEnum(str, Enum):
    FOO = 'foo'
    BAR = 'bar'


class MyModel(BaseModel):
    x: Literal[MyEnum.FOO]


MyModel(x=MyEnum.FOO)  # Enum instance, works
# MyModel(x=MyEnum.BAR)  # Fails, as expected
MyModel(x='foo')       # String value, works