Fix GenericPrefetch invariance with QuerySet (#3380)

Author: pmdevita

django-stubs/contrib/contenttypes/prefetch.pyi

@@ -1,3 +1,4 @@
+from collections.abc import Sequence
 from typing import Any, Generic
 
 from django.db.models import Model, Prefetch
@@ -9,7 +10,7 @@ from typing_extensions import TypeVar, override
 _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]]
+    "_PrefetchedQuerySetsT", bound=Sequence[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

django-stubs/db/models/query.pyi

@@ -303,7 +303,7 @@ class Prefetch(Generic[_LookupT, _PrefetchedQuerySetT, _ToAttrT]):
     def add_prefix(self, prefix: str) -> None: ...
     def get_current_prefetch_to(self, level: int) -> str: ...
     def get_current_to_attr(self, level: int) -> tuple[str, str]: ...
-    def get_current_querysets(self, level: int) -> list[_PrefetchedQuerySetT] | None: ...
+    def get_current_querysets(self, level: int) -> Sequence[_PrefetchedQuerySetT] | None: ...
 
 def normalize_prefetch_lookups(lookups: Sequence[str | Prefetch], prefix: str | None = None) -> list[Prefetch]: ...
 def prefetch_related_objects(model_instances: Sequence[_Model], *related_lookups: str | Prefetch) -> None: ...

tests/typecheck/managers/querysets/test_prefetch_related.yml

@@ -428,7 +428,7 @@
         - myapp
     main: |
         from django.contrib.contenttypes.prefetch import GenericPrefetch
-        from myapp.models import Bookmark, Animal, TaggedItem
+        from myapp.models import Bookmark, Animal, TaggedItem, Dog, Cat
         from typing_extensions import reveal_type
 
         # Basic GenericPrefetch usage
@@ -437,6 +437,17 @@
         )
         reveal_type(prefetch)  # N: Revealed type is "django.contrib.contenttypes.prefetch.GenericPrefetch[Literal['content_object'], list[django.db.models.query.QuerySet[django.db.models.base.Model, django.db.models.base.Model]], str]"
 
+        # GenericPrefetch usage with custom querysets
+        prefetch_custom = GenericPrefetch(
+            "content_object", [Dog.objects.all(), Cat.objects.all()]
+        )
+        reveal_type(prefetch_custom)  # N: Revealed type is "django.contrib.contenttypes.prefetch.GenericPrefetch[Literal['content_object'], list[myapp.models.CustomQuerySet[django.db.models.base.Model, django.db.models.base.Model]], str]"
+
+        prefetch_custom_mixed = GenericPrefetch(
+            "content_object", [Bookmark.objects.all(), Animal.objects.only("name"), Dog.objects.all(), Cat.objects.all()]
+        )
+        reveal_type(prefetch_custom_mixed)  # N: Revealed type is "django.contrib.contenttypes.prefetch.GenericPrefetch[Literal['content_object'], list[django.db.models.query.QuerySet[django.db.models.base.Model, django.db.models.base.Model]], 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]"
@@ -477,6 +488,17 @@
                 class Animal(models.Model):
                     name = models.CharField(max_length=100)
 
+                class CustomQuerySet(models.QuerySet):
+                    pass
+
+                CustomManager = models.Manager.from_queryset(CustomQuerySet)
+
+                class Dog(models.Model):
+                    objects = CustomManager()
+
+                class Cat(models.Model):
+                    objects = CustomManager()
+
                 class TaggedItem(models.Model):
                     tag = models.CharField(max_length=100)
                     content_type = models.ForeignKey(ContentType, on_delete=models.CASCADE)