Fix Self typed custom queryset methods incompatible with base queryse…

…t type (#1840)
typeddjango · Nov 24, 2023 · 1b0aeb6 · 1b0aeb6
1 parent 0a61d81
commit 1b0aeb6
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 15 deletions.
diff --git 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
@@ -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
@@ -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: