How to specify properties of certain objects via typing?

user_4045379 0
Viewed 34

I have a certain database model in my Django app that has a nullable FK field that is strictly related to its other field, like so:

class ModelType(models.TextChoices):
    non_empty = "non-empty"
    empty = "empty"


class SomeModel(models.Model):
    model_type = models.CharField(choices=ModelType.choices)
    related_item = models.ForeignKey(to=[...], null=True)

    class Meta:
        constraints = [
            models.CheckConstraint(
                check=(
                    Q(model_type="empty", related_item__isnull=True)
                    | Q(model_type="non-empty", related_item__isnull=False)
                ),
                name="related_item_required_only_when_non_empty",
        ]

I would like to be able to write a function that returns queryset of SomeModel objects with certain properties (for typing purposes):

def get_non_empty_models() -> QuerySet[NonEmptySomeModel]:
    return SomeModel.objects.filter(model_type=ModelType.non_empty)

>>> print([
...    some_model.related_item.func() 
...    for some_model in get_non_empty_models()
... ])  # MyPy complains about related_item being None

Is there a way of annotating such qualities of objects for MyPy?

mypy django-models django python
1 Answers

This cannot be automatically inferred without plugin support, but you can manually acknowledge mypy the type of related_item with a TypeVar by making SomeModel generic.

T = TypeVar("T")

class SomeModel(models.Model, Generic[T]):
    model_type = models.CharField(choices=ModelType.choices)
    related_item: T = models.ForeignKey(to=[...], null=True)

    class Meta:
        constraints = [
            models.CheckConstraint(
                check=(
                    Q(model_type="empty", related_item__isnull=True)
                    | Q(model_type="non-empty", related_item__isnull=False)
                ),
                name="related_item_required_only_when_non_empty",
        ]

and annotate the helper functions this way:

def get_non_empty_models() -> QuerySet[SomeModel[RelatedItem]]:
    ...

def get_empty_models() -> QuerySet[SomeModel[Optional[RelatedItem]]]:
    ...

If you are already using third-party stubs like django-stubs, it's likely that you will need to use cast to narrow the type in ways like cast(QuerySet[SomeModel[RelatedItem]], something).

user_116123990
Related