From 1b0aeb6e24d5f2808ab9705ab3974ff455808396 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 | 6 ++--- .../managers/querysets/test_from_queryset.yml | 4 ++-- 3 files changed, 17 insertions(+), 15 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..d107e82b1f 100644 --- a/tests/typecheck/managers/querysets/test_as_manager.yml +++ b/tests/typecheck/managers/querysets/test_as_manager.yml @@ -1,10 +1,10 @@ - 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]]" + 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]]" installed_apps: - myapp 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: