Last active
June 23, 2023 14:36
-
-
Save edinella/cd7ad3a8e4c62fb5b23504174fa110e8 to your computer and use it in GitHub Desktop.
How Pydantic can handle `id` as string and `_id` as ObjectId for PyMongo
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from typing import List, Optional | |
from bson import ObjectId | |
from pydantic import BaseModel, Field | |
class PyObjectId(ObjectId): | |
@classmethod | |
def __get_validators__(cls): | |
yield cls.validate | |
@classmethod | |
def validate(cls, v): | |
if not ObjectId.is_valid(v): | |
raise ValueError("Invalid objectid") | |
return ObjectId(v) | |
@classmethod | |
def __modify_schema__(cls, field_schema): | |
field_schema.update(type="string") | |
class Entity(BaseModel): | |
id: Optional[PyObjectId] = Field(alias="_id") | |
ref: Optional[PyObjectId] = Field() | |
refs: Optional[List[PyObjectId]] = Field() | |
class Config: | |
allow_population_by_field_name = True # Pydantic should use the alias when populating the model from a dictionary (which allows you to pass in a dictionary with an _id key rather than an id key) | |
json_encoders = {ObjectId: str} # ObjectId should be encoded to a string when converting the model to JSON | |
test = Entity( | |
id="64760b13500e13df454ff03d", | |
ref="647e0fa0b4ee216624761493", | |
refs=["64760b13500e13df454ff03d", "647e0fa0b4ee216624761493"], | |
) | |
print(test) | |
# id=ObjectId('64760b13500e13df454ff03d') ref=ObjectId('647e0fa0b4ee216624761493') refs=[ObjectId('64760b13500e13df454ff03d'), ObjectId('647e0fa0b4ee216624761493')] | |
print(test.dict(by_alias=True)) # ensure that the dictionary uses the field aliases (so it has an _id key rather than an id key) | |
# {'_id': ObjectId('64760b13500e13df454ff03d'), 'ref': ObjectId('647e0fa0b4ee216624761493'), 'refs': [ObjectId('64760b13500e13df454ff03d'), ObjectId('647e0fa0b4ee216624761493')]} | |
print(test.json()) | |
# {"id": "64760b13500e13df454ff03d", "ref": "647e0fa0b4ee216624761493", "refs": ["64760b13500e13df454ff03d", "647e0fa0b4ee216624761493"]} | |
print(test.id) | |
# 64760b13500e13df454ff03d |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment