CODEWITHSUNDEEP
pythondjangopython-typingdjango-stubs
Johnny Metz
Johnny Metz

Reputation: 6025

Type Hint for Django model with annotated field

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.

Upvotes: 6

Views: 4374

Answers (2)

Niel Godfrey P. Ponciano
Niel Godfrey P. Ponciano

Reputation: 10709

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]]:
    ...

Upvotes: -1

JGC
JGC

Reputation: 6373

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.

Upvotes: -1

Related Questions