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.

Move

Answer 1

Python 3.9 introduces typing.Annotated e.g. typing.Annotated[Toolbox, "WithTools"].

So for scenarios where we just need to indicate it for clarity that the object is annotated:

from typing import Annotated

from django.db.models import QuerySet

def get_toolbox_to_small_tools_mappings(
    toolboxes: QuerySet | list[Annotated[Toolbox, "WithTools"]],
) -> dict[Annotated[Toolbox, "WithTools"], list[Tool]]:
    ...

Apyretic

But this doesn't actually solve anything right? You still have to use # type-ignore? – Mcneese

Answer 2

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.

Pass

I disagree. Refactor tools won't translate this kind of code. Would be better to use (even if buggy) type annotations and defer "fixing" it until typehints are first class citizens – Cessation

@boatcoder, I agree and I hadn't really thought about the impact on refactoring tools. I will update my answer to say it's not a good approach. – Pass

What refactor tools do you have in mind? Why won't they translate this kind of code? – Mcneese

In the example we're referring to the small_tools attribute using a string "small_tools" via get_attr. If I ever wanted to rename the small_tools attribute on the Toolbox class I could use vscode or PyCharm refactor tools and the tool goes through the code base and renames all uses of small_tools to my new name. By using a string here the refactor tools won't detect this as a use of that attribute (it just sees a string), so will leave it alone, thus breaking the code. – Pass