Reparametrize implicit generic QuerySet subclasses (#3217)

When a custom QuerySet subclass is defined without explicit type
parameters (e.g. `class MyQS(QuerySet): ...`), make it implicitly
generic by copying the parent QuerySet's type variables, analogous
to the existing `reparametrize_any_manager_hook` for Manager
subclasses.

This enables the annotate plugin hook to propagate annotation type
information through custom querysets via `copy_modified(args=...)`,
which requires the queryset class to have type variables.

Author: federicobond

mypy_django_plugin/main.py

@@ -43,6 +43,7 @@
     add_as_manager_to_queryset_class,
     create_new_manager_class_from_from_queryset_method,
     reparametrize_any_manager_hook,
+    reparametrize_any_queryset_hook,
     resolve_manager_method,
 )
 from mypy_django_plugin.transformers.models import (
@@ -232,6 +233,8 @@ def get_customize_class_mro_hook(self, fullname: str) -> Callable[[ClassDefConte
         info = self._get_typeinfo_or_none(fullname)
         if info and info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME):
             return reparametrize_any_manager_hook
+        if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
+            return reparametrize_any_queryset_hook
         return None
 
     @override

mypy_django_plugin/transformers/managers.py

@@ -589,6 +589,58 @@ def _defer() -> None:
     )
 
 
+def reparametrize_any_queryset_hook(ctx: ClassDefContext) -> None:
+    """
+    Add implicit generics to QuerySet subclasses that are defined without generic.
+
+    Eg.
+
+        class MyQuerySet(models.QuerySet): ...
+
+    is interpreted as:
+
+        _Model = TypeVar('_Model', bound=Model, covariant=True)
+        _Row = TypeVar('_Row', covariant=True, default=_Model)
+        class MyQuerySet(models.QuerySet[_Model, _Row]): ...
+
+    Note that this does not happen if mypy is run with disallow_any_generics = True,
+    as not specifying the generic type is then considered an error.
+    """
+    queryset = ctx.api.lookup_fully_qualified_or_none(ctx.cls.fullname)
+    if queryset is None or queryset.node is None:
+        return
+    assert isinstance(queryset.node, TypeInfo)
+
+    if queryset.node.type_vars:
+        # We've already been here
+        return
+
+    parent_queryset = next(
+        (base for base in queryset.node.bases if base.type.has_base(fullnames.QUERYSET_CLASS_FULLNAME)),
+        None,
+    )
+    if parent_queryset is None or len(parent_queryset.args) != 2:
+        return
+
+    model_param = get_proper_type(parent_queryset.args[0])
+    if not isinstance(model_param, AnyType) or model_param.type_of_any is not TypeOfAny.from_omitted_generics:
+        return
+
+    type_vars = tuple(parent_queryset.type.defn.type_vars)
+
+    # If we end up with placeholders we need to defer so the placeholders are
+    # resolved in a future iteration
+    if any(has_placeholder(type_var) for type_var in type_vars):
+        if not ctx.api.final_iteration:
+            ctx.api.defer()
+        else:
+            return
+
+    parent_queryset.args = type_vars
+    queryset.node.defn.type_vars = list(type_vars)
+    queryset.node.add_type_vars()
+
+
 def reparametrize_any_manager_hook(ctx: ClassDefContext) -> None:
     """
     Add implicit generics to manager classes that are defined without generic.

tests/typecheck/fields/test_related.yml

@@ -718,9 +718,9 @@
         reveal_type(user.article_set) # N: Revealed type is "myapp.models.Article_RelatedManager"
         reveal_type(user.book_set.add)  # N: Revealed type is "def (*objs: myapp.models.Book | builtins.int, bulk: builtins.bool =)"
         reveal_type(user.article_set.add)  # N: Revealed type is "def (*objs: myapp.models.Article | builtins.int, bulk: builtins.bool =)"
-        reveal_type(user.book_set.filter)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.LibraryEntityQuerySet"
+        reveal_type(user.book_set.filter)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.LibraryEntityQuerySet[myapp.models.Book, myapp.models.Book]"
         reveal_type(user.book_set.get())  # N: Revealed type is "myapp.models.Book"
-        reveal_type(user.article_set.filter)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.LibraryEntityQuerySet"
+        reveal_type(user.article_set.filter)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.LibraryEntityQuerySet[myapp.models.Article, myapp.models.Article]"
         reveal_type(user.article_set.get())  # N: Revealed type is "myapp.models.Article"
         reveal_type(user.book_set.queryset_method())  # N: Revealed type is "builtins.int"
         reveal_type(user.article_set.queryset_method())  # N: Revealed type is "builtins.int"

tests/typecheck/managers/querysets/test_annotate.yml

@@ -552,3 +552,35 @@
         from django.db import models
         class MyModel(models.Model):
             name = models.CharField(max_length=100)
+
+- case: annotate_on_implicitly_generic_custom_queryset
+  main: |
+    from typing_extensions import reveal_type
+    from myapp.models import MyModel
+    from django.db.models import Count
+
+    qs = MyModel.objects.all().annotate(num_items=Count("id"))
+    obj = qs.get()
+    reveal_type(obj.num_items)  # N: Revealed type is "Any"
+    obj.nonexistent  # E: "MyModel@AnnotatedWith[TypedDict({'num_items': Any})]" has no attribute "nonexistent"  [attr-defined]
+
+    # Custom queryset methods remain available after annotate
+    qs.custom_method()
+  installed_apps:
+    - myapp
+  files:
+    - path: myapp/__init__.py
+    - path: myapp/models.py
+      content: |
+        from django.db import models
+        from typing_extensions import Self
+
+        class MyQuerySet(models.QuerySet):
+            def custom_method(self) -> Self:
+                return self.filter(name="test")
+
+        MyManager = models.Manager.from_queryset(MyQuerySet)
+
+        class MyModel(models.Model):
+            name = models.CharField(max_length=100)
+            objects = MyManager()

tests/typecheck/managers/querysets/test_from_queryset.yml

@@ -1042,3 +1042,32 @@
                 # Forward-referenced metaclass defined after use
                 class ForwardMCS(type):
                     pass
+
+-   case: subclass_queryset_without_type_parameters_disallow_any_generics
+    main: |
+        from typing_extensions import reveal_type
+        from myapp.models import MyModel
+        reveal_type(MyModel.objects)
+        reveal_type(MyModel.objects.get())
+    installed_apps:
+        - myapp
+    mypy_config: |
+        [mypy-myapp.models]
+        disallow_any_generics = true
+    out: |
+        main:3: note: Revealed type is "myapp.models.ManagerFromMyQuerySet[myapp.models.MyModel]"
+        main:4: note: Revealed type is "myapp.models.MyModel"
+        myapp/models:3: error: Missing type parameters for generic type "QuerySet"  [type-arg]
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+
+                class MyQuerySet(models.QuerySet):
+                    pass
+
+                MyManager = models.Manager.from_queryset(MyQuerySet)
+
+                class MyModel(models.Model):
+                    objects = MyManager()

tests/typecheck/managers/querysets/test_union_type.yml

@@ -10,7 +10,7 @@
 
         reveal_type(model_cls) # N: Revealed type is "type[myapp.models.Order] | type[myapp.models.User]"
         reveal_type(model_cls.objects)  # N: Revealed type is "myapp.models.ManagerFromMyQuerySet[myapp.models.Order] | myapp.models.ManagerFromMyQuerySet[myapp.models.User]"
-        reveal_type(model_cls.objects.my_method())  # N: Revealed type is "myapp.models.MyQuerySet"
+        reveal_type(model_cls.objects.my_method())  # N: Revealed type is "myapp.models.MyQuerySet[Any, Any]"
     installed_apps:
         - myapp
     files: