From eaa27011b2f96db008379181d21f03df39705575 Mon Sep 17 00:00:00 2001 From: Federico Bond Date: Tue, 24 Mar 2026 15:38:22 +1100 Subject: [PATCH] Reparametrize implicit generic QuerySet subclasses 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. --- mypy_django_plugin/main.py | 3 ++ mypy_django_plugin/transformers/managers.py | 52 +++++++++++++++++++ tests/typecheck/fields/test_related.yml | 4 +- .../managers/querysets/test_annotate.yml | 32 ++++++++++++ .../managers/querysets/test_from_queryset.yml | 29 +++++++++++ .../managers/querysets/test_union_type.yml | 2 +- 6 files changed, 119 insertions(+), 3 deletions(-) diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py index 05ef0f361..0fbb01928 100644 --- a/mypy_django_plugin/main.py +++ b/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 diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index 066584046..8ba746196 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/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. diff --git a/tests/typecheck/fields/test_related.yml b/tests/typecheck/fields/test_related.yml index 213953079..ae13ed234 100644 --- a/tests/typecheck/fields/test_related.yml +++ b/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" diff --git a/tests/typecheck/managers/querysets/test_annotate.yml b/tests/typecheck/managers/querysets/test_annotate.yml index e024652f4..9ec219dc3 100644 --- a/tests/typecheck/managers/querysets/test_annotate.yml +++ b/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() diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index 64b09caa7..fd5f1e279 100644 --- a/tests/typecheck/managers/querysets/test_from_queryset.yml +++ b/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() diff --git a/tests/typecheck/managers/querysets/test_union_type.yml b/tests/typecheck/managers/querysets/test_union_type.yml index e361e8fe4..5c126bba0 100644 --- a/tests/typecheck/managers/querysets/test_union_type.yml +++ b/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: