Как указать свойства определенных объектов через ввод?

У меня есть определенная модель базы данных в моем приложении Django, которая имеет нулевое поле FK, которое строго связано с другим полем, примерно так:

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",
        ]

Я хотел бы иметь возможность написать функцию, которая возвращает queryset из SomeModel объектов с определенными свойствами (для целей типизации):

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

Есть ли способ аннотировать такие качества объектов для MyPy?

Это не может быть автоматически выведено без поддержки плагина, но вы можете вручную подтвердить mypy тип related_item с помощью TypeVar, сделав 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",
        ]

и аннотируйте вспомогательные функции таким образом:

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

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

Если вы уже используете сторонние заглушки типа django-stubs, то, скорее всего, вам понадобится использовать cast для сужения типа такими способами, как cast(QuerySet[SomeModel[RelatedItem]], something).

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