Skip to content

Instantly share code, notes, and snippets.

@CartfordSam
Last active March 19, 2024 16:31
Show Gist options
  • Save CartfordSam/3eaa30864a3199ee6ca2484fc4f4fa15 to your computer and use it in GitHub Desktop.
Save CartfordSam/3eaa30864a3199ee6ca2484fc4f4fa15 to your computer and use it in GitHub Desktop.
Pydantic v1.10 handle _id with extended BaseModel
from pydantic import BaseModel, root_validator
# NOTE: this works (for me) with pydantic==1.10.13,
# but we'll be upgrading to v2 and you can do the same
# thing with the model_validator
# considerations:
# - I'm using this in FastAPI, mostly read-only to my users,
# so I don't have a need to keep the ObjectId, and convert
# back to the bson ObjectId as-needed
class BaselineOverrideID(BaseModel):
"""Inherit this to convert the bson.ObjectID to a string, and
set the variable name to return using the _id_str_name variable
USAGE:
class TestOverride(BaselineOverrideID):
_id_str_name = "test_id"
test_id: Optional[str] = "gets_over_written"
...
As long as the dict passed into the constructor for this has an _id,
it overwrites the test_id default value with the str-equivalent.
EXAMPLE:
test_dict = {"_id": bson.ObjectId("64a3752f7946675e994894df")}
test_obj = TestOverride.parse_obj(test_dict)
test_obj.test_id
>>> "64a3752f7946675e994894df"
type(test_obj.test_id)
>>> str
"""
_id_str_name: str
# We have to set pre=True since the inheriting objects
# likely don't keep the _id field
@root_validator(pre=True)
def _set_fields(cls, values: dict) -> dict:
"""This is a validator that converts the _id
field into a string with the specified name.
Args:
values (dict): Stores the attributes of the extended object.
"""
if "_id" in values.keys():
values[cls._id_str_name] = str(values["_id"])
return values
if __name__=="__main__":
import bson
class ExampleUser(BaselineOverrideID):
_id_str_name = "user_id"
user_id: Optional[str] = None
email: str
# ... + more fields
example_user_d = {
"_id": bson.ObjectId("63a3752f7946675e994894df"),
"email": "sam@uptrends.ai"
}
example_no_id_user_d = {
"email": "idk@internet.xyz"
}
example_user = ExampleUser(**example_user_d)
example_no_id_user = ExampleUser(**example_no_id_user_d)
print(f"user w/ id: type(user_id)=={type(example_user.user_id)} ({example_user.user_id})")
print(f"user w/o id: type(user_id)=={type(example_no_id_user.user_id)} ({example_no_id_user.user_id})")
"""OUTPUT >>
user w/ id: type(user_id)==<class 'str'> (63a3752f7946675e994894df)
user w/o id: type(user_id)==<class 'NoneType'> (None)
"""
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment