Pydantic/SQLAlchemy: How to work with enums?

What is the best way to convert a sqlalchemy model to a pydantic schema (model) if it includes an enum field?

Sqlalchemy

import enum
from sqlalchemy import Enum, Column, String
from sqlalchemy.orm import declarative_base

Base = declarative_base()

class StateEnum(enum.Enum):
    CREATED = 'CREATED'
    UPDATED = 'UPDATED'

class Adapter(Base):
    __tablename__ = 'adapters'
    id = Column(String, primary_key=True)
    friendly_name = Column(String(256), nullable=False)
    state: StateEnum = Column(Enum(StateEnum))

Pydantic

from pydantic import BaseModel
from enum import Enum

class StateEnumDTO(str, Enum):
    CREATED = 'CREATED'
    UPDATED = 'UPDATED'


class AdapterDTO(BaseModel):
    friendly_name: str
    state: StateEnumDTO  # This currently cannot be converted?

    class Config:
        allow_population_by_field_name = True
        orm_mode = True
        use_enum_values = True

Conversion

AdapterDTO.from_orm(Adapter(friendly_name='test', state=StateEnum.CREATED))

This leads to the error

value is not a valid enumeration member; permitted: 'CREATED', 'UPDATED' (type=type_error.enum; enum_values=[<StateEnumDTO.CREATED: 'CREATED'>, <StateEnumDTO.UPDATED: 'UPDATED'>])

How can I configure either

a.) the serialization with the from_orm method?

b.) the creation of the state field?

c.) How to convert it the other way around?

Is there a native way to do this with pydantic or how is this typically done?

Update: Test case

def test_enum_conversion_to_dto():
    adapter = Adapter(id='1', friendly_name='test', state=StateEnum.CREATED)
    adapter_dto = AdapterDTO.from_orm(adapter)
    assert adapter_dto.state == StateEnumDTO.CREATED
    assert adapter_dto.state.value == StateEnum.CREATED.value

Answer

Pydantic requires that both enum classes have the same type definition.

In your case, StateEnum inherits from enum.Enum, but StateEnumDTO inherits from both str and enum.Enum.

You can fix this issue by changing your SQLAlchemy enum definition:

class StateEnum(str, enum.Enum):
    CREATED = 'CREATED'
    UPDATED = 'UPDATED'

Advertisement

Answer