How do I make a generic reference type in mypy that also accepts own subclasses?

Viewed 74

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.

0 Answers
Related