CODEWITHSUNDEEP
pythonjsonserializationpydantic
Vu Viet Dung
Vu Viet Dung

Reputation: 383

How to serialize custom type that extends builtin type in pydantic?

currently I'm working with FastAPI and pydantic as serializer. Problem is, we're using snowflake id on the server side, which means we need to convert those ids to string before sending to client (javascript) because the id is larger than JS's MAX SAFE INTEGER.

So I tried to create a new class which extends python's int type and customize how it will be serialized and deserialized. Here's my code:

class SnowflakeId(int):
    @classmethod
    def __get_validators__(cls):
        yield cls.validate

    @classmethod
    def validate(cls, v: str):
        return int(v)

    @classmethod
    def __modify_schema__(cls, field_schema: dict) -> None:
        field_schema['type'] = 'string'

And here is the model:

class BaseModel(pydantic.BaseModel):
    __abstract__ = True

    id: SnowflakeId

    class Config:
        orm_mode = True
        arbitrary_types_allowed = True
        json_encoders = {
            SnowflakeId: lambda v: str(v)
        }
        alias_generator = camelize
        allow_population_by_field_name = True

It works fine when deserializing from json string into int id, however, when it comes to the serialization, the output still is integer. I want it to serialize the id into string also, is it possible?

Upvotes: 9

Views: 6722

Answers (1)

kosciej16
kosciej16

Reputation: 7158

Yes it is!

json_encoders is a good try, however under the hood pydantic calls json.dumps. So for serializable types (like your SnowflakeId) it won't care about additional json_encoders.

What you can do is to override dumps method:

def my_dumps(v, *, default):
    for key, value in v.items():
        if isinstance(value, SnowflakeId):
            v[key] = str(value)
        else:
            v[key] = value
    return json.dumps(v)

class BaseModel(pydantic.BaseModel):
    id: SnowflakeId

    class Config:
        json_dumps = my_dumps

And let validate return SnowflakeId:

class SnowflakeId(int):
    ...

    @classmethod
    def validate(cls, v: str):
        return cls(v)

m = BaseModel(id="123")
print(m.json()) # {"id": "123"}

Upvotes: 4

Related Questions