Подсказка типа для модели Django с аннотированным полем

Допустим, у меня есть следующие модели Django:

class Toolbox(models.Model):
    name = models.CharField(max_length=255)
    tools = models.ManyToManyField("Tool")

class Tool(models.Model):
    class Size(models.TextChoices):
        SMALL = "S"
        MEDIUM = "M"
        LARGE = "L"

    name = models.CharField(max_length=255)
    size = models.CharField(max_length=10, choices=Size.choices)

У меня есть функция для получения всех маленьких инструментов для каждого набора инструментов. Подсказка типа аргумента исходит от этого SO ответа :

from django.db.models import QuerySet

def get_toolbox_to_small_tools_mappings(
    toolboxes: QuerySet | list[Toolbox],
) -> dict[Toolbox, list[Tool]]:
    return {toolbox: toolbox.small_tools for toolbox in toolboxes}

Идея здесь состоит в том, чтобы потребовать от пользователей этой функции предварительной выборки этого поля small_tools с помощью prefetch_related(), чтобы уменьшить количество обращений к БД и ускорить код:

toolboxes = Toolbox.objects.prefetch_related(
    Prefetch(
        "tools",
        queryset=Tool.objects.filter(size=Tool.Size.SMALL),
        to_attr="small_tools",
    )
)
toolbox_to_small_tools_mappings = get_toolbox_to_small_tools_mappings(toolboxes)

Все это работает отлично, но mypy жалуется со следующей ошибкой:

Ошибка: "Toolbox" не имеет атрибута "small_tools" [attr-defined]

Есть ли способ исправить это?

Тип WithAnnotations[Model] из django-subs ( см. здесь) является вариантом, но он багги.

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