Support GenericPrefetch (#2851)

Author: UnknownPlatypus

django-stubs/contrib/contenttypes/prefetch.pyi

@@ -1,9 +1,20 @@
-from typing import Any
+from typing import Any, Generic, TypeVar
 
-from django.db.models import Prefetch
+from django.db.models import Model, Prefetch
 from django.db.models.query import QuerySet
 
-class GenericPrefetch(Prefetch):
-    def __init__(self, lookup: str, querysets: list[QuerySet], to_attr: str | None = None) -> None: ...
+# The type of the lookup passed to Prefetch(...)
+# This will be specialized to a `LiteralString` in the plugin for further processing and validation
+_LookupT = TypeVar("_LookupT", bound=str, covariant=True)
+# The type of the querysets passed to GenericPrefetch(...)
+_PrefetchedQuerySetsT = TypeVar(
+    "_PrefetchedQuerySetsT", bound=list[QuerySet[Model]], covariant=True, default=list[QuerySet[Model]]
+)
+# The attribute name to store the prefetched list[_PrefetchedQuerySetT]
+# This will be specialized to a `LiteralString` in the plugin for further processing and validation
+_ToAttrT = TypeVar("_ToAttrT", bound=str, covariant=True, default=str)
+
+class GenericPrefetch(Prefetch, Generic[_LookupT, _PrefetchedQuerySetsT, _ToAttrT]):
+    def __init__(self, lookup: _LookupT, querysets: _PrefetchedQuerySetsT, to_attr: _ToAttrT | None = None) -> None: ...
     def __getstate__(self) -> dict[str, Any]: ...
-    def get_current_querysets(self, level: int) -> list[QuerySet] | None: ...
+    def get_current_querysets(self, level: int) -> _PrefetchedQuerySetsT | None: ...

mypy_django_plugin/lib/fullnames.py

@@ -17,6 +17,7 @@
 MANAGER_CLASS_FULLNAME = "django.db.models.manager.Manager"
 RELATED_MANAGER_CLASS = "django.db.models.fields.related_descriptors.RelatedManager"
 PREFETCH_CLASS_FULLNAME = "django.db.models.query.Prefetch"
+GENERIC_PREFETCH_CLASS_FULLNAME = "django.contrib.contenttypes.prefetch.GenericPrefetch"
 
 CHOICES_CLASS_FULLNAME = "django.db.models.enums.Choices"
 CHOICES_TYPE_METACLASS_FULLNAME = "django.db.models.enums.ChoicesType"

mypy_django_plugin/transformers/querysets.py

@@ -512,7 +512,12 @@ def check_valid_attr_value(
 
 
 def check_valid_prefetch_related_lookup(
-    ctx: MethodContext, lookup: str, django_model: DjangoModel, django_context: DjangoContext
+    ctx: MethodContext,
+    lookup: str,
+    django_model: DjangoModel,
+    django_context: DjangoContext,
+    *,
+    is_generic_prefetch: bool = False,
 ) -> bool:
     """Check if a lookup string resolve to something that can be prefetched"""
     current_model_cls = django_model.cls
@@ -528,7 +533,17 @@ def check_valid_prefetch_related_lookup(
                 ctx.context,
             )
             return False
-        if isinstance(rel_obj_descriptor, ForwardManyToOneDescriptor):
+        if contenttypes_installed and is_generic_prefetch:
+            from django.contrib.contenttypes.fields import GenericForeignKey
+
+            if not isinstance(rel_obj_descriptor, GenericForeignKey):
+                ctx.api.fail(
+                    f'"{through_attr}" on "{current_model_cls.__name__}" is not a GenericForeignKey, '
+                    f"GenericPrefetch can only be used with GenericForeignKey fields",
+                    ctx.context,
+                )
+                return True
+        elif isinstance(rel_obj_descriptor, ForwardManyToOneDescriptor):
             current_model_cls = rel_obj_descriptor.field.remote_field.model
         elif isinstance(rel_obj_descriptor, ReverseOneToOneDescriptor):
             current_model_cls = rel_obj_descriptor.related.related_model  # type:ignore[assignment] # Can't be 'self' for non abstract models
@@ -563,7 +578,10 @@ def check_valid_prefetch_related_lookup(
 
 
 def check_conflicting_lookups(
-    ctx: MethodContext, observed_attr: str, qs_types: dict[str, Instance | None], queryset_type: Instance | None
+    ctx: MethodContext,
+    observed_attr: str,
+    qs_types: dict[str, Instance | None],
+    queryset_type: Instance | None,
 ) -> bool:
     is_conflicting_lookup = bool(observed_attr in qs_types and qs_types[observed_attr] != queryset_type)
     if is_conflicting_lookup:
@@ -641,7 +659,13 @@ def extract_prefetch_related_annotations(ctx: MethodContext, django_context: Dja
             )
             qs_types[to_attr] = queryset_type
         if not to_attr and lookup:
-            check_valid_prefetch_related_lookup(ctx, lookup, qs_model, django_context)
+            check_valid_prefetch_related_lookup(
+                ctx,
+                lookup,
+                qs_model,
+                django_context,
+                is_generic_prefetch=typ.type.has_base(fullnames.GENERIC_PREFETCH_CLASS_FULLNAME),
+            )
             check_conflicting_lookups(ctx, lookup, qs_types, queryset_type)
             qs_types[lookup] = queryset_type
 

tests/typecheck/managers/querysets/test_prefetch_related.yml

@@ -384,3 +384,99 @@
                     subject_content_type = models.ForeignKey(ContentType, on_delete=models.CASCADE)
                     subject_id = models.PositiveIntegerField()
                     subject = GenericForeignKey("subject_content_type", "subject_id")
+
+-   case: django_contrib_contenttypes_generic_prefetch
+    installed_apps:
+        - django.contrib.contenttypes
+        - myapp
+    main: |
+        from django.contrib.contenttypes.prefetch import GenericPrefetch
+        from myapp.models import Bookmark, Animal, TaggedItem
+        from typing_extensions import reveal_type
+
+        # Basic GenericPrefetch usage
+        prefetch = GenericPrefetch(
+            "content_object", [Bookmark.objects.all(), Animal.objects.only("name")]
+        )
+        reveal_type(prefetch)  # N: Revealed type is "django.contrib.contenttypes.prefetch.GenericPrefetch[Literal['content_object'], builtins.list[django.db.models.query.QuerySet[django.db.models.base.Model, django.db.models.base.Model]], builtins.str]"
+
+        # Using GenericPrefetch with prefetch_related
+        qs = TaggedItem.objects.prefetch_related(prefetch).all()
+        reveal_type(qs)  # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.TaggedItem, myapp.models.TaggedItem]"
+
+        # GenericPrefetch with to_attr
+        prefetch_with_attr = GenericPrefetch(
+            "content_object",
+            [Bookmark.objects.all(), Animal.objects.only("name")],
+            to_attr="prefetched_object"
+        )
+        qs_with_attr = TaggedItem.objects.prefetch_related(prefetch_with_attr).all()
+        reveal_type(qs_with_attr)  # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.TaggedItem@AnnotatedWith[TypedDict({'prefetched_object': builtins.list[Any]})], myapp.models.TaggedItem@AnnotatedWith[TypedDict({'prefetched_object': builtins.list[Any]})]]"
+        reveal_type(qs_with_attr.get().prefetched_object)  # N: Revealed type is "builtins.list[Any]"
+
+        # GenericPrefetch on invalid field (not a GenericForeignKey)
+        regular_fk_prefetch = GenericPrefetch(
+            "content_type",
+            [Bookmark.objects.all(), Animal.objects.only("name")],
+        )
+        TaggedItem.objects.prefetch_related(regular_fk_prefetch).all() # E: "content_type" on "TaggedItem" is not a GenericForeignKey, GenericPrefetch can only be used with GenericForeignKey fields  [misc]
+        regular_field_prefetch = GenericPrefetch(
+            "tag",
+            [Bookmark.objects.all(), Animal.objects.only("name")],
+        )
+        TaggedItem.objects.prefetch_related(regular_field_prefetch).all() # E: "tag" on "TaggedItem" is not a GenericForeignKey, GenericPrefetch can only be used with GenericForeignKey fields  [misc]
+
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                from django.contrib.contenttypes.models import ContentType
+                from django.contrib.contenttypes.fields import GenericForeignKey
+
+                class Bookmark(models.Model):
+                    url = models.URLField()
+
+                class Animal(models.Model):
+                    name = models.CharField(max_length=100)
+
+                class TaggedItem(models.Model):
+                    tag = models.CharField(max_length=100)
+                    content_type = models.ForeignKey(ContentType, on_delete=models.CASCADE)
+                    object_id = models.PositiveIntegerField()
+                    content_object = GenericForeignKey('content_type', 'object_id')
+
+
+-   case: uninstalled_django_contrib_contenttypes_generic_prefetch
+    installed_apps:
+        - myapp
+    main: |
+        from django.contrib.contenttypes.prefetch import GenericPrefetch
+        from myapp.models import Bookmark, Animal, TaggedItem
+        from typing_extensions import reveal_type
+
+        # Basic GenericPrefetch usage
+        prefetch = GenericPrefetch(
+            "content_object", [Bookmark.objects.all(), Animal.objects.only("name")]
+        )
+        reveal_type(prefetch)  # N: Revealed type is "django.contrib.contenttypes.prefetch.GenericPrefetch[Literal['content_object'], builtins.list[django.db.models.query.QuerySet[django.db.models.base.Model, django.db.models.base.Model]], builtins.str]"
+
+        # Using GenericPrefetch with prefetch_related
+        qs = TaggedItem.objects.prefetch_related(prefetch).all() # E: Cannot find "content_object" on "TaggedItem" object, "content_object" is an invalid parameter to "prefetch_related()"  [misc]
+        reveal_type(qs)  # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.TaggedItem, myapp.models.TaggedItem]"
+
+
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+
+                class Bookmark(models.Model):
+                    url = models.URLField()
+
+                class Animal(models.Model):
+                    name = models.CharField(max_length=100)
+
+                class TaggedItem(models.Model):
+                    tag = models.CharField(max_length=100)