I am using Pydantic on top of MongoDB, and things are working pretty well including documentation for the IDs I'm describing here. However, I want to be able to type references between different documents as well. I was able to make it work halfway as can be seen in the final part of this snippet.
from typing import TypeVar, Generic
from pydantic import BaseModel, Field
ModelType = TypeVar("ModelType", bound="MyBaseModel")
class IDType(str):
pass
class ReferenceId(IDType, Generic[ModelType]):
pass
class MyBaseModel(BaseModel):
id: IDType = Field(default_factory=IDType)
class MySubclassModel(MyBaseModel):
pass
class AnotherSubclassModel(MyBaseModel):
reference_to_my_subclass: ReferenceId[MySubclassModel]
class YetAnotherSubclassModel(MyBaseModel):
reference_to_another_subclass: ReferenceId[AnotherSubclassModel]
# Obviously fails but I want those to pass:
a = AnotherSubclassModel(reference_to_my_subclass=MyBaseModel().id)
# Expected type 'ReferenceId[MySubclassModel]', got 'IDType' instead
x = YetAnotherSubclassModel(reference_to_another_subclass=a.id)
# Expected type 'ReferenceId[AnotherSubclassModel]', got 'IDType' instead
# But this works as I want
def my_function(reference_to_my_subclass: ReferenceId[MySubclassModel]):
pass
my_function(a.reference_to_my_subclass)
# While these will fail, and I want them to fail:
# This one fails for the reason that I don't want, but at least fails
my_function(a.id)
# Expected type 'ReferenceId[MySubclassModel]', got 'IDType' instead
# This one for the right reason
my_function(x.reference_to_another_subclass)
# Expected type 'ReferenceId[MySubclassModel]', got 'ReferenceId[AnotherSubclassModel]' instead
It might be that the solution lies in somehow having MyBaseModel.id reference "this subclass class", but I am not sure how to accomplish that.