Declare Prefetch as generic and do specialization in the plugin (#2786)

Author: UnknownPlatypus

django-stubs/db/models/query.pyi

@@ -15,6 +15,12 @@ _Model = TypeVar("_Model", bound=Model, covariant=True)
 _Row = TypeVar("_Row", covariant=True, default=_Model)  # ONLY use together with _Model
 _TupleT = TypeVar("_TupleT", bound=tuple[Any, ...], covariant=True)
 
+# The type of the prefetched model
+_PrefetchedQuerySetT = TypeVar("_PrefetchedQuerySetT", bound=QuerySet[Model], covariant=True, default=QuerySet[Model])
+# The attribute name to use 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)
+
 MAX_GET_RESULTS: int
 REPR_OUTPUT_SIZE: int
 
@@ -234,21 +240,23 @@ class RawQuerySet(Iterable[_Model], Sized):
 # Deprecated alias of QuerySet, for compatibility only.
 _QuerySet: TypeAlias = QuerySet  # noqa: PYI047
 
-class Prefetch:
+class Prefetch(Generic[_PrefetchedQuerySetT, _ToAttrT]):
     prefetch_through: str
     prefetch_to: str
-    queryset: QuerySet | None
-    to_attr: str | None
-    def __init__(self, lookup: str, queryset: QuerySet | None = None, to_attr: str | None = None) -> None: ...
+    queryset: _PrefetchedQuerySetT | None
+    to_attr: _ToAttrT | None
+    def __init__(
+        self, lookup: str, queryset: _PrefetchedQuerySetT | None = None, to_attr: _ToAttrT | None = None
+    ) -> None: ...
     def __getstate__(self) -> dict[str, Any]: ...
     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]: ...
     @deprecated(
         "get_current_queryset() is deprecated and will be removed in Django 6.0. Use get_current_querysets() instead."
     )
-    def get_current_queryset(self, level: int) -> QuerySet | None: ...
-    def get_current_querysets(self, level: int) -> list[QuerySet] | None: ...
+    def get_current_queryset(self, level: int) -> _PrefetchedQuerySetT | None: ...
+    def get_current_querysets(self, level: int) -> list[_PrefetchedQuerySetT] | None: ...
 
 def prefetch_related_objects(model_instances: Iterable[_Model], *related_lookups: str | Prefetch) -> None: ...
 async def aprefetch_related_objects(model_instances: Iterable[_Model], *related_lookups: str | Prefetch) -> None: ...

ext/django_stubs_ext/patch.py

@@ -22,7 +22,7 @@
 from django.db.models.lookups import Lookup
 from django.db.models.manager import BaseManager
 from django.db.models.options import Options
-from django.db.models.query import BaseIterable, ModelIterable, QuerySet, RawQuerySet
+from django.db.models.query import BaseIterable, ModelIterable, Prefetch, QuerySet, RawQuerySet
 from django.forms.formsets import BaseFormSet
 from django.forms.models import BaseModelForm, BaseModelFormSet, ModelChoiceField, ModelFormOptions
 from django.utils.connection import BaseConnectionHandler, ConnectionProxy
@@ -97,6 +97,7 @@ def __repr__(self) -> str:
     MPGeneric(BaseIterable),
     MPGeneric(ForwardManyToOneDescriptor),
     MPGeneric(ReverseOneToOneDescriptor),
+    MPGeneric(Prefetch),
 ]
 
 

mypy_django_plugin/main.py

@@ -143,6 +143,9 @@ def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], MypyTy
             if info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME):
                 return querysets.determine_proper_manager_type
 
+            if info.has_base(fullnames.PREFETCH_CLASS_FULLNAME):
+                return partial(querysets.specialize_prefetch_type, django_context=self.django_context)
+
         return None
 
     @cached_property

mypy_django_plugin/transformers/querysets.py

@@ -7,7 +7,7 @@
 from mypy.checker import TypeChecker
 from mypy.nodes import ARG_NAMED, ARG_NAMED_OPT, CallExpr, Expression
 from mypy.plugin import FunctionContext, MethodContext
-from mypy.types import AnyType, Instance, TupleType, TypedDictType, TypeOfAny, get_proper_type
+from mypy.types import AnyType, Instance, LiteralType, TupleType, TypedDictType, TypeOfAny, get_proper_type
 from mypy.types import Type as MypyType
 
 from mypy_django_plugin.django.context import DjangoContext, LookupsAreUnsupported
@@ -318,20 +318,37 @@ def extract_proper_type_queryset_values(ctx: MethodContext, django_context: Djan
     return default_return_type.copy_modified(args=[django_model.typ, row_type])
 
 
-def _infer_prefetch_element_model_type(queryset_expr: Expression | None, api: TypeChecker) -> Instance | None:
+def _infer_prefetch_queryset_type(queryset_expr: Expression, api: TypeChecker) -> Instance | None:
     """Infer the model Instance from `Prefetch(queryset=...)`"""
-    if queryset_expr is None:
-        # TODO: Infer the model type from the lookup in `Prefetch(lookup=..., to_attr=...)`
-        return None
     try:
         qs_type = get_proper_type(api.expr_checker.accept(queryset_expr))
     except Exception:
         return None
     if isinstance(qs_type, Instance):
-        return helpers.extract_model_type_from_queryset(qs_type, api)
+        return qs_type
     return None
 
 
+def specialize_prefetch_type(ctx: FunctionContext, django_context: DjangoContext) -> MypyType:
+    """Function hook for `Prefetch(...)` to specialize its `to_attr` generic parameters."""
+    default = get_proper_type(ctx.default_return_type)
+    if not isinstance(default, Instance):
+        return ctx.default_return_type
+
+    api = helpers.get_typechecker_api(ctx)
+
+    # Guaranteed to exist because the TypeVar has a default
+    to_attr_type = default.args[1]
+    # We specialize the `to_attr` str type to a Literal[str] so that it can be used later
+    # to annotate the correct attribute name on the model instances and do further validations
+    # See `extract_prefetch_related_annotations` below.
+    if to_attr_expr := helpers.get_call_argument_by_name(ctx, "to_attr"):
+        if to_attr_value := helpers.resolve_string_attribute_value(to_attr_expr, django_context):
+            to_attr_type = LiteralType(value=to_attr_value, fallback=api.named_generic_type("builtins.str", []))
+
+    return default.copy_modified(args=[default.args[0], to_attr_type])
+
+
 def extract_prefetch_related_annotations(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
     """
     Extract annotated attributes via `prefetch_related(Prefetch(..., to_attr=...))`
@@ -353,18 +370,57 @@ def extract_prefetch_related_annotations(ctx: MethodContext, django_context: Dja
 
     for expr, typ in zip(ctx.args[0], ctx.arg_types[0], strict=False):
         typ = get_proper_type(typ)
-        if not (
-            isinstance(typ, Instance)
-            and typ.type.fullname == fullnames.PREFETCH_CLASS_FULLNAME
-            and isinstance(expr, CallExpr)
-            and (to_attr_expr := helpers.get_class_init_argument_by_name(expr, "to_attr"))
-            and (to_attr_value := helpers.resolve_string_attribute_value(to_attr_expr, django_context))
-        ):
+        if not (isinstance(typ, Instance) and typ.type.has_base(fullnames.PREFETCH_CLASS_FULLNAME)):
             continue
 
-        # Determine model type from the `queryset` attr
-        queryset_expr = helpers.get_class_init_argument_by_name(expr, "queryset")
-        elem_model = _infer_prefetch_element_model_type(queryset_expr, api)
+        # 1) Try to get to_attr from specialized type arg
+        to_attr_value: str | None = None
+        if (
+            len(typ.args) >= 2
+            and isinstance((to_attr_t := get_proper_type(typ.args[1])), LiteralType)
+            and isinstance(to_attr_t.value, str)
+        ):
+            to_attr_value = to_attr_t.value
+
+        # Fallback: parse inline call expression
+        if to_attr_value is None:
+            if not (
+                isinstance(expr, CallExpr)
+                and (to_attr_expr := helpers.get_class_init_argument_by_name(expr, "to_attr"))
+                and (to_attr_value := helpers.resolve_string_attribute_value(to_attr_expr, django_context))
+            ):
+                continue
+
+        # 2) Determine element model type from specialized type arg
+        elem_model: Instance | None = None
+        if len(typ.args) >= 1 and isinstance((queryset_type := get_proper_type(typ.args[0])), Instance):
+            elem_model = helpers.extract_model_type_from_queryset(queryset_type, api)
+
+        # Fallback: parse inline call expression
+        if elem_model is None and isinstance(expr, CallExpr):
+            queryset_expr = helpers.get_class_init_argument_by_name(expr, "queryset")
+            if queryset_expr is not None and isinstance(
+                (inferred_queryset_type := _infer_prefetch_queryset_type(queryset_expr, api)), Instance
+            ):
+                elem_model = helpers.extract_model_type_from_queryset(inferred_queryset_type, api)
+            else:
+                # Resolve model type using the "lookup" required first argument and
+                # the model associated with the current queryset.
+                lookup_expr = helpers.get_class_init_argument_by_name(expr, "lookup")
+                if lookup_expr is None:
+                    continue
+
+                if lookup_value := helpers.resolve_string_attribute_value(lookup_expr, django_context):
+                    if django_model := helpers.get_model_info_from_qs_ctx(ctx, django_context):
+                        try:
+                            observed_model_cls = django_context.resolve_lookup_into_field(
+                                django_model.cls, lookup_value
+                            )[1]
+                            if model_info := helpers.lookup_class_typeinfo(api, observed_model_cls):
+                                elem_model = Instance(model_info, [])
+                        except (FieldError, LookupsAreUnsupported):
+                            pass
+
         value_type = api.named_generic_type(
             "builtins.list",
             [elem_model if elem_model is not None else AnyType(TypeOfAny.special_form)],

tests/typecheck/managers/querysets/test_prefetch_related.yml

@@ -1,12 +1,28 @@
 -   case: prefetch_related_to_attr
     main: |
         from myapp.models import Article, Tag
-        from django.db.models import Prefetch
+        from django.db.models import Prefetch, F, QuerySet
         from django.db import models
+        from typing_extensions import Literal, TypedDict
+        from django_stubs_ext import WithAnnotations
 
-        # Noop (to_attr not provided)
+        ### Noop (to_attr not provided)
         reveal_type(Article.objects.prefetch_related(Prefetch("tags")).all()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Article, myapp.models.Article]"
         reveal_type(Article.objects.prefetch_related(Prefetch("tags", Tag.objects.all())).all()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Article, myapp.models.Article]"
+        reveal_type(Prefetch("tags")) # N: Revealed type is "django.db.models.query.Prefetch[django.db.models.query.QuerySet[django.db.models.base.Model, django.db.models.base.Model], builtins.str]"
+        reveal_type(Prefetch("tags", Tag.objects.all())) # N: Revealed type is "django.db.models.query.Prefetch[django.db.models.query.QuerySet[myapp.models.Tag, myapp.models.Tag], builtins.str]"
+
+        # Prefetch created in a function with no to_attr
+        def get_prefetch_no_to_attr() -> Prefetch[QuerySet[Tag]]:
+            return Prefetch("tags", Tag.objects.all())
+
+        reveal_type(Article.objects.prefetch_related(get_prefetch_no_to_attr()).all()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Article, myapp.models.Article]"
+
+        # Prefetch created in a function with no to_attr and no queryset
+        def get_prefetch_no_to_attr_no_qs() -> Prefetch:
+            return Prefetch("tags")
+
+        reveal_type(Article.objects.prefetch_related(get_prefetch_no_to_attr_no_qs()).all()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Article, myapp.models.Article]"
 
         # On the QuerySet
         article_qs = Article.objects.all().prefetch_related(Prefetch("tags", Tag.objects.all(), to_attr="every_tags"))
@@ -59,29 +75,54 @@
         )
         reveal_type(mixed_plain) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Article@AnnotatedWith[TypedDict({'arts3': builtins.list[myapp.models.Article]})], myapp.models.Article@AnnotatedWith[TypedDict({'arts3': builtins.list[myapp.models.Article]})]]"
 
-
-        ## Not Supported
-
-        # Prefetch with `to_attr` arg but without the `queryset` arg
-        # TODO: We should be able to resolve a more accurate type using existing lookup `resolve_lookup_expected_type` machinery
-        reveal_type(Article.objects.prefetch_related(models.Prefetch("tags", to_attr="just_tags")).get().just_tags) # N: Revealed type is "builtins.list[Any]"
-
         # Intermediary variable -- function scope
         def foo() -> None:
             tag_prefetch = Prefetch("tags", Tag.objects.all(), to_attr="every_tags")
-            reveal_type(Article.objects.prefetch_related(tag_prefetch).all()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Article, myapp.models.Article]"
+            reveal_type(Article.objects.prefetch_related(tag_prefetch).all()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Article@AnnotatedWith[TypedDict({'every_tags': builtins.list[myapp.models.Tag]})], myapp.models.Article@AnnotatedWith[TypedDict({'every_tags': builtins.list[myapp.models.Tag]})]]"
 
         # Intermediary variable -- module scope
         tag_prefetch = Prefetch("tags", Tag.objects.all(), to_attr="every_tags")
-        reveal_type(Article.objects.prefetch_related(tag_prefetch).all()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Article, myapp.models.Article]"
+        reveal_type(Article.objects.prefetch_related(tag_prefetch).all()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Article@AnnotatedWith[TypedDict({'every_tags': builtins.list[myapp.models.Tag]})], myapp.models.Article@AnnotatedWith[TypedDict({'every_tags': builtins.list[myapp.models.Tag]})]]"
+
+        # Prefetch created in a function
+        def get_invalid_prefetch() -> Prefetch[QuerySet[Tag], Literal["every_tags"]]:
+            return Prefetch("tags", Tag.objects.all(), to_attr="foo") # E: Incompatible return value type (got "Prefetch[QuerySet[Tag, Tag], Literal['foo']]", expected "Prefetch[QuerySet[Tag, Tag], Literal['every_tags']]")  [return-value] # E: Argument "to_attr" to "Prefetch" has incompatible type "Literal['foo']"; expected "Literal['every_tags'] | None"  [arg-type]
+
+        def get_prefetch() -> Prefetch[QuerySet[Tag], Literal["every_tags"]]:
+            return Prefetch("tags", Tag.objects.all(), to_attr="every_tags")
+
+        reveal_type(Article.objects.prefetch_related(get_prefetch()).all()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Article@AnnotatedWith[TypedDict({'every_tags': builtins.list[myapp.models.Tag]})], myapp.models.Article@AnnotatedWith[TypedDict({'every_tags': builtins.list[myapp.models.Tag]})]]"
+
+        # Prefetch created in a function with to_attr as an intermediary variable
+        def get_prefetch_with_var() -> Prefetch[QuerySet[Tag], Literal["every_tags"]]:
+            to_attr = "every_tags" # TODO: RM error next line
+            return Prefetch("tags", Tag.objects.all(), to_attr) # E: Argument 3 to "Prefetch" has incompatible type "str"; expected "Literal['every_tags'] | None"  [arg-type]
+
+        reveal_type(Article.objects.prefetch_related(get_prefetch_with_var()).all()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Article@AnnotatedWith[TypedDict({'every_tags': builtins.list[myapp.models.Tag]})], myapp.models.Article@AnnotatedWith[TypedDict({'every_tags': builtins.list[myapp.models.Tag]})]]"
 
         # Mixed inline `Prefetch` and variable `Prefetch` in one call
         mixed_qs = Article.objects.prefetch_related(
             tag_prefetch,
             Prefetch("article_set", Article.objects.all(), to_attr="arts2"),
         )
-        reveal_type(mixed_qs) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Article@AnnotatedWith[TypedDict({'arts2': builtins.list[myapp.models.Article]})], myapp.models.Article@AnnotatedWith[TypedDict({'arts2': builtins.list[myapp.models.Article]})]]"
+        reveal_type(mixed_qs) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Article@AnnotatedWith[TypedDict({'every_tags': builtins.list[myapp.models.Tag], 'arts2': builtins.list[myapp.models.Article]})], myapp.models.Article@AnnotatedWith[TypedDict({'every_tags': builtins.list[myapp.models.Tag], 'arts2': builtins.list[myapp.models.Article]})]]"
+
+        # Prefetch with `to_attr` arg but without the `queryset` arg
+        reveal_type(Article.objects.prefetch_related(models.Prefetch("tags", to_attr="just_tags")).get().just_tags) # N: Revealed type is "builtins.list[myapp.models.Tag]"
+
+        # Prefetch with annotated `queryset`
+        reveal_type(Article.objects.prefetch_related( # N: Revealed type is "builtins.list[myapp.models.Tag@AnnotatedWith[TypedDict({'foo': Any})]]"
+            models.Prefetch("tags", Tag.objects.annotate(foo=F("id")).all(), to_attr="just_tags")).get().just_tags
+        )
+
+        class MyDict(TypedDict):
+            foo: str
+
+        def get_prefetch_with_annotations() -> Prefetch[QuerySet[WithAnnotations[Tag, MyDict]], Literal["every_tags"]]:
+            return Prefetch("tags", Tag.objects.annotate(foo=F("id")), "every_tags")
 
+        reveal_type(Article.objects.prefetch_related(get_prefetch_with_annotations()).all()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Article@AnnotatedWith[TypedDict({'every_tags': builtins.list[myapp.models.Tag@AnnotatedWith[TypedDict('main.MyDict', {'foo': builtins.str})]]})], myapp.models.Article@AnnotatedWith[TypedDict({'every_tags': builtins.list[myapp.models.Tag@AnnotatedWith[TypedDict('main.MyDict', {'foo': builtins.str})]]})]]"
+        reveal_type(Article.objects.prefetch_related(get_prefetch_with_annotations()).get().every_tags[0].foo) # N: Revealed type is "builtins.str"
     installed_apps:
         - myapp
     files: