Type Hint for Django model with annotated field

Question

Let's say I have the following Django models:

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)

I have a function to get all small tools for each toolbox. The argument type hint comes from this SO answer:

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}

The idea here is to require users of this function to prefetch this small_tools field using prefetch_related() to reduce the number of db hits and speed up the code:

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)

This all works great but mypy is complaining with the following error:

error: "Toolbox" has no attribute "small_tools" [attr-defined]

Is there anyway to fix this?

The WithAnnotations[Model] type from django-subs (see here) is an option but it's buggy.

Answer 1

Your best options may be to use type annotations from django-stubs or disable annotations (type: ignore).

One alternative option is to switch to using getattr, but this has down-sides (see below).

from django.db.models import QuerySet

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

This keeps things simple, and maintains the same functionality, but at the cost of affecting refactoring tools and other developer IDE functionality that won't be able to translate this code.