[Python] Mixin for making third-party Django model fields typed

Given the problem that not all libraries add type annotations, I am trying to build a mixin that would allow me to add type hints easily. Without that, the type-checker infers the type of the model instance field as that of the underlying built-in field’s value (e.g., str, int, date, etc.) that the custom field is built upon (or Unknown).

Here is an example:

from django.db import models

# ========== defined in a third-party library ==========
class StructuredCountry:
iso_code: str
...

class CountryField(models.CharField):
# Not shown: the textual value from the database
# is converted to a StructuredCountry object.
pass

# ========== defined in my own code ====================
class MyModel(models.Model):
country = CountryField()

reveal_type(MyModel().country) # ⇒ str
MyModel().country.iso_code # E: Attribute "iso_code" is unknown


For the above example, I would like to be able to:

class TypedCountryField(FieldAnnotationMixin, CountryField):
pass

class MyModel(models.Model):
country = TypedCountryField()

reveal_type(MyModel().country) # ⇒ StructuredCountry
MyModel().country.iso_code # ✔


So far I have come up with the following (inspired by VSCode’s stubs), which indeed works.

if TYPE_CHECKING:
class FieldAnnotationMixin[FieldClassType: type[models.Field], ValueType]:
def __new__(cls, *args, **kwargs) -> FieldClassType: ...

# Class access
@overload
def __get__(self, instance: None, owner: Any) -> Self: ... # type: ignore[overload]

# Model instance access
@overload
def __get__(self, instance: models.Model, owner: Any) -> ValueType: ...

# Non-Model instance access
@overload
def __get__(self, instance: Any, owner: Any) -> Self: ...
else:
class FieldAnnotationMixin[IgnoredT1, IgnoredT2]:
pass


However, I am stuck on the nullable field bit. For a field defined on the model with null=True, the model instance value should be ValueType | None (or Optional[ValueType]) and the type-checker will complain if the value is used before verifying that it is not null. Looking again at VSCode’s stubs I know that I can overload the __new__ method:

@overload
def __new__(cls, *args, null: Literal[False] = False, **kwargs) -> FieldClassType[ValueType]: ...

@overload
def __new__(cls, *args, null: Literal[True] = True, **kwargs) -> FieldClassType[ValueType | None]: ...


I am running into two problems here:

1. The type-checker (rightly) complains “TypeVar "type[FieldClassType@FieldAnnotationMixin]" is not subscriptable”. And I don’t know in advance whether the third-party field is generic or not.
2. I have no idea how to transfer the knowledge of the null parameter to the __get__ overload so that it will return Optional[ValueType].

Continue reading...