Last active
March 19, 2024 16:31
-
-
Save CartfordSam/3eaa30864a3199ee6ca2484fc4f4fa15 to your computer and use it in GitHub Desktop.
Pydantic v1.10 handle _id with extended BaseModel
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 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