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].

Вернуться на верх