CODEWITHSUNDEEP
pythonpydantic
Florian Brucker
Florian Brucker

Reputation: 10365

Initialize a Literal enum in a pydantic model

pydantic supports regular enums just fine, and one can initialize an enum-typed field using both an enum instance and an enum 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 also supports 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

Now I want to combine enums and literals, i.e. force a field value to equal one particular enum instance. However, I still want to be able to pass the correct enum value (i.e. not only the correct enum instance), and that doesn't seem to work:

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

Upvotes: 15

Views: 13339

Answers (2)

rocksteady
rocksteady

Reputation: 2550

You might have noticed that an Enum based class attribute is not of type str. Accessing the actual value of the attribute can be achieved using .value.

Please refer to: https://docs.python.org/3/library/enum.html

Example:

>>> from enum import Enum

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

>>> MyEnum.FOO
<MyEnum.FOO: 'foo'>
>>> type(MyEnum.FOO)
<enum 'MyENum'>

You need to retrieve the actual value of it like this:

>>> MyEnum.FOO.value
'foo'
>>> type(MyEnum.FOO.value)
str
>>> MyEnum.FOO.name
'FOO'

Upvotes: 2

alex_noname
alex_noname

Reputation: 32233

Try to use string based Enum. Like so:

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

Upvotes: 20

Related Questions