From e18a070885f2a7109d3646f4024556f6f5071ee9 Mon Sep 17 00:00:00 2001 From: moran abadie Date: Fri, 24 Nov 2023 18:34:48 +0100 Subject: [PATCH] Fix Self typed custom queryset methods incompatible with base queryset type (#1840) --- mypy_django_plugin/transformers/managers.py | 22 ++++++++++--------- .../managers/querysets/test_as_manager.yml | 21 ++++++++++++++---- .../managers/querysets/test_from_queryset.yml | 4 ++-- 3 files changed, 31 insertions(+), 16 deletions(-) diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index 1baf96dc1a..c1d0e91779 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -109,19 +109,21 @@ def _process_dynamic_method( variables = method_type.variables ret_type = method_type.ret_type + if not is_fallback_queryset: + queryset_instance = Instance(queryset_info, manager_instance.args) + else: + # The fallback queryset inherits _QuerySet, which has two generics + # instead of the one exposed on QuerySet. That means that we need + # to add the model twice. In real code it's not possible to inherit + # from _QuerySet, as it doesn't exist at runtime, so this fix is + # only needed for plugin-generated querysets. + queryset_instance = Instance(queryset_info, [manager_instance.args[0], manager_instance.args[0]]) + # For methods on the manager that return a queryset we need to override the # return type to be the actual queryset class, not the base QuerySet that's # used by the typing stubs. if method_name in MANAGER_METHODS_RETURNING_QUERYSET: - if not is_fallback_queryset: - ret_type = Instance(queryset_info, manager_instance.args) - else: - # The fallback queryset inherits _QuerySet, which has two generics - # instead of the one exposed on QuerySet. That means that we need - # to add the model twice. In real code it's not possible to inherit - # from _QuerySet, as it doesn't exist at runtime, so this fix is - # only needed for pluign-generated querysets. - ret_type = Instance(queryset_info, [manager_instance.args[0], manager_instance.args[0]]) + ret_type = queryset_instance variables = [] args_types = method_type.arg_types[1:] if _has_compatible_type_vars(base_that_has_method): @@ -138,7 +140,7 @@ def _process_dynamic_method( ] if base_that_has_method.self_type: # Manages -> Self returns - ret_type = _replace_type_var(ret_type, base_that_has_method.self_type.fullname, manager_instance) + ret_type = _replace_type_var(ret_type, base_that_has_method.self_type.fullname, queryset_instance) # Drop any 'self' argument as our manager is already initialized return method_type.copy_modified( diff --git a/tests/typecheck/managers/querysets/test_as_manager.yml b/tests/typecheck/managers/querysets/test_as_manager.yml index 9ef966acad..8345b0bd4b 100644 --- a/tests/typecheck/managers/querysets/test_as_manager.yml +++ b/tests/typecheck/managers/querysets/test_as_manager.yml @@ -1,10 +1,11 @@ - case: self_return_management main: | - from myapp.models import MyModel - reveal_type(MyModel.objects.example_simple()) # N: Revealed type is "myapp.models.ManagerFromMyQuerySet[myapp.models.MyModel]" - reveal_type(MyModel.objects.example_list()) # N: Revealed type is "builtins.list[myapp.models.ManagerFromMyQuerySet[myapp.models.MyModel]]" + from myapp.models import MyModel, MyModel2 + reveal_type(MyModel.objects.example_simple()) # N: Revealed type is "myapp.models.MyQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects.example_list()) # N: Revealed type is "builtins.list[myapp.models.MyQuerySet[myapp.models.MyModel]]" reveal_type(MyModel.objects.example_simple().just_int()) # N: Revealed type is "builtins.int" - reveal_type(MyModel.objects.example_dict()) # N: Revealed type is "builtins.dict[builtins.str, myapp.models.ManagerFromMyQuerySet[myapp.models.MyModel]]" + reveal_type(MyModel.objects.example_dict()) # N: Revealed type is "builtins.dict[builtins.str, myapp.models.MyQuerySet[myapp.models.MyModel]]" + reveal_type(MyModel2.objects.method()) # N: Revealed type is "myapp.models.MyQuerySet2" installed_apps: - myapp @@ -19,6 +20,10 @@ class BaseQuerySet(models.QuerySet): def example_dict(self) -> Dict[str, Self]: ... + class MyQuerySet2(models.QuerySet["MyModel2"]): + def method(self) -> "MyQuerySet2": + return self + class MyQuerySet(BaseQuerySet): def example_simple(self) -> Self: ... def example_list(self) -> List[Self]: ... @@ -26,6 +31,9 @@ class MyModel(models.Model): objects = MyQuerySet.as_manager() + + class MyModel2(models.Model): + objects = MyQuerySet2.as_manager() - case: declares_manager_type_like_django main: | from myapp.models import MyModel @@ -192,6 +200,8 @@ reveal_type(MyOtherModel.objects.dummy_override()) # N: Revealed type is "myapp.models.MyOtherModel" reveal_type(MyOtherModel.objects.example_mixin(MyOtherModel())) # N: Revealed type is "myapp.models.MyOtherModel" reveal_type(MyOtherModel.objects.example_other_mixin()) # N: Revealed type is "myapp.models.MyOtherModel" + reveal_type(MyOtherModel.objects.test_self()) # N: Revealed type is "myapp.models._MyModelQuerySet2[myapp.models.MyOtherModel]" + reveal_type(MyOtherModel.objects.test_sub_self()) # N: Revealed type is "myapp.models._MyModelQuerySet2[myapp.models.MyOtherModel]" installed_apps: - myapp files: @@ -200,6 +210,7 @@ content: | from typing import TypeVar, Generic from django.db import models + from typing_extensions import Self T = TypeVar("T", bound=models.Model) T_2 = TypeVar("T_2", bound=models.Model) @@ -215,12 +226,14 @@ def override(self) -> T: ... def override2(self) -> T: ... def dummy_override(self) -> int: ... + def test_sub_self(self) -> Self: ... class _MyModelQuerySet2(SomeMixin, _MyModelQuerySet[T_2]): def example_2(self) -> T_2: ... def override(self) -> T_2: ... def override2(self) -> T_2: ... def dummy_override(self) -> T_2: ... # type: ignore[override] + def test_self(self) -> Self: ... class MyModelQuerySet(_MyModelQuerySet2["MyModel"]): def override(self) -> "MyModel": ... diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index 0c8f1f5621..65c126c5cd 100644 --- a/tests/typecheck/managers/querysets/test_from_queryset.yml +++ b/tests/typecheck/managers/querysets/test_from_queryset.yml @@ -1,8 +1,8 @@ - case: from_queryset_self_return_management main: | from myapp.models import MyModel - reveal_type(MyModel.objects.example_simple()) # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]" - reveal_type(MyModel.objects.example_list()) # N: Revealed type is "builtins.list[myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]]" + reveal_type(MyModel.objects.example_simple()) # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects.example_list()) # N: Revealed type is "builtins.list[myapp.models.ModelQuerySet[myapp.models.MyModel]]" installed_apps: - myapp files: