pythonsqlalchemypydantic

How to make pydantic not serialize custom enum to dict


I created this piece of code to simulate the issue. Let's say we have a custom enum that is an enum of states and has two values. It's full name and short version:

from dataclasses import dataclass, field
from enum import Enum
from typing import Any

from pydantic import BaseModel, model_serializer


@dataclass
class StateDataMixin:
    state: str
    short: str = field(repr=False)

    def __repr__(self) -> str:
        return self.state

    def __hash__(self):
        return hash(self.state)


class State(StateDataMixin, Enum):
    ALABAMA = "Alabama", "AL"
    ALASKA = "Alaska", "AK"
    ...

    @property
    def value(self):
        return self._value_.state  # type: ignore


class SomeState(BaseModel):
    state: State


some_state = SomeState(state=State.ALASKA)
result = some_state.model_dump()

I have a pydantic object, that I would like to serialise into the database. The model goes into the function like this one:

async def create(
    self, db: AsyncSession, *, obj_in: ItemCreate
) -> Appointment:
    db_obj = self.model(**obj_in.model_dump())
    db.add(db_obj)
    db.commit()
    return db_obj

That's why in the simulated code above I used:

some_state = SomeState(state=State.ALASKA)
result = some_state.model_dump()

Apparently this serialises into: {'state': {'state': 'Alaska', 'short': 'AK'}}

And SQLAlchemy throws an error when it gets this type of serialised object:

sqlalchemy.exc.StatementError: (builtins.TypeError) unhashable type: 'dict'

My SQLalchemy model looks something like this:

class Item(Base):
    __tablename__ = "items"

    id: Mapped[int]
    title: Mapped[str]
    state: Mapped[State]

Where State refers to mentioned enum. From the SQLAlchemy side everything works fine. When I added:

def __repr__(self) -> str:
    return self.state

def __hash__(self):
    return hash(self.state)

to the StateDataMixin class SQLAlchemy was able to properly create database model based on this enum and understands it correctly.

How can I make Pydantic properly serialise my enum so it can be understood by the SQLAlchemy? I guess what I'm trying to achieve is that instead of:
{'state': {'state': 'Alaska', 'short': 'AK'}}
I would like to have
{'state': State.ALASKA}
I suppose then this should be understood correctly by SQLAlchemy.


Solution

  • I could make this work with this Pydantic model (the trick, I suspect, is in returning the enum's name):

    class ItemModel(BaseModel):
        model_config = ConfigDict(use_enum_values=False, from_attributes=True)
        state: State
        title: str
    
        @field_serializer('state')
        def serialize_state(self, state: State, _info):
            print(f'Returning {state.state=}')
            return state.name
    

    with which this SQLAlchemy insertion round-trips successfully:

    item_model = ItemModel(title='A', state=State.ALASKA)
    
    with Session.begin() as s:
        item = Item(**item_model.model_dump())
        s.add(item)
    
    with Session() as s:
        item = s.scalar(sa.select(Item).where(Item.title == 'A').limit(1))
        print(item.title, item.state)