Storing enums

[aldoshkind]

Hi!
I faced problem with serializing enums: I get bson.errors.InvalidDocument: cannot encode object: <State.Preparation: 'Preparation'>, of type: <enum 'State'> when trying to store enum.

What is the correct solution?

Full example code:

from enum import Enum
from typing import Optional
from uuid import UUID, uuid4

from pydantic import BaseModel
from pydantic_mongo import AbstractRepository
from pymongo import MongoClient

client = MongoClient("localhost:27017", uuidRepresentation='standard')
database = client["pydantic_mongo_db"]
db = database["test_enum"]

class State(Enum):
    Preparation = 'Preparation'
    Processing = 'Processing'

class Info(BaseModel):
    id: UUID = uuid4()
    state: Optional[State] = State('Preparation')

class InfoRepository(AbstractRepository[ Info ]):
   class Meta:
      collection_name = 'info_repository'

repo = InfoRepository(db)

info = Info()

repo.save(info)

[Artucuno]

Hey there!

Here is the solution to your issue.

from enum import Enum
from typing import Optional

from pydantic import BaseModel
from pymongo import MongoClient

from pydantic_mongo import AbstractRepository, ObjectIdField

client = MongoClient("localhost:27017")
database = client["test_enum"]


class State(str, Enum):  # <---
    Preparation = 'Preparation'
    Processing = 'Processing'


class Info(BaseModel):
    id: ObjectIdField = None  # <---
    state: Optional[State] = State.Preparation


class InfoRepository(AbstractRepository[Info]):
    class Meta:
        collection_name = 'info_repository'


repo = InfoRepository(database=database)

info = Info()

repo.save(info)

res = [x for x in repo.find_by({})]
print(res)
for f in res:
    print(type(State(f.state)))