CODEWITHSUNDEEP
pythonsqlalchemypydantic
Mateusz Anikiej
Mateusz Anikiej

Reputation: 108

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.

Upvotes: 0

Views: 292

Answers (1)

snakecharmerb
snakecharmerb

Reputation: 55864

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)

Upvotes: 1

Related Questions