Fix Self typed custom queryset methods incompatible with base queryset type (#1840) (#1852)

Author: moranabadie

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(

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, MyModelWithoutSelf
+        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(MyModelWithoutSelf.objects.method())  # N: Revealed type is "myapp.models.QuerySetWithoutSelf"
 
     installed_apps:
         - myapp
@@ -26,6 +27,13 @@
 
                 class MyModel(models.Model):
                     objects = MyQuerySet.as_manager()
+
+                class QuerySetWithoutSelf(models.QuerySet["MyModelWithoutSelf"]):
+                    def method(self) -> "QuerySetWithoutSelf":
+                        return self
+
+                class MyModelWithoutSelf(models.Model):
+                    objects = QuerySetWithoutSelf.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": ...

tests/typecheck/managers/querysets/test_from_queryset.yml

@@ -1,8 +1,12 @@
 -   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]]"
+        from myapp.models import MyModel, MyModelWithoutSelf
+        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.MyQuerySet[myapp.models.MyModel]]"
+        reveal_type(MyModel.objects.test_custom_manager())  # N: Revealed type is "myapp.models.CustomManagerFromMyQuerySet[myapp.models.MyModel]"
+        reveal_type(MyModelWithoutSelf.objects.method())  # N: Revealed type is "myapp.models.QuerySetWithoutSelf"
     installed_apps:
         - myapp
     files:
@@ -11,16 +15,34 @@
             content: |
                 from django.db import models
                 from django.db.models.manager import BaseManager
+                from typing import List, Dict
                 from typing_extensions import Self
-                from typing import List
 
-                class ModelQuerySet(models.QuerySet):
+                class CustomManager(BaseManager):
+                    def test_custom_manager(self) -> Self: ...
+
+                class BaseQuerySet(models.QuerySet):
+                    def example_dict(self) -> Dict[str, Self]: ...
+
+                class MyQuerySet(BaseQuerySet):
                     def example_simple(self) -> Self: ...
                     def example_list(self) -> List[Self]: ...
-                NewManager = BaseManager.from_queryset(ModelQuerySet)
+                    def just_int(self) -> int: ...
+
+                NewManager = CustomManager.from_queryset(MyQuerySet)
+
                 class MyModel(models.Model):
                     objects = NewManager()
 
+                class QuerySetWithoutSelf(models.QuerySet["MyModelWithoutSelf"]):
+                    def method(self) -> "QuerySetWithoutSelf":
+                        return self
+
+                ManagerWithoutSelf = BaseManager.from_queryset(QuerySetWithoutSelf)
+
+                class MyModelWithoutSelf(models.Model):
+                    objects = ManagerWithoutSelf()
+
 -   case: from_queryset_with_base_manager
     main: |
         from myapp.models import MyModel