fix as manager and from queryset self types (#1788) (#1789)

typeddjango · Oct 21, 2023 · 0ce571e · 0ce571e
1 parent 1334efd
commit 0ce571e
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 0 deletions.
diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py
@@ -132,6 +132,9 @@ def _process_dynamic_method(
             _replace_type_var(arg_type, base_that_has_method.defn.type_vars[0].fullname, manager_instance.args[0])
             for arg_type in args_types
         ]
+    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)
 
     # 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,3 +1,31 @@
+-   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().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]]"
+
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                from typing import List, Dict
+                from typing_extensions import 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]: ...
+                    def just_int(self) -> int: ...
+
+                class MyModel(models.Model):
+                    objects = MyQuerySet.as_manager()
 -   case: declares_manager_type_like_django
     main: |
         from myapp.models import MyModel

diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml
@@ -1,3 +1,26 @@
+-   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]]"
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                from django.db.models.manager import BaseManager
+                from typing_extensions import Self
+                from typing import List
+
+                class ModelQuerySet(models.QuerySet):
+                    def example_simple(self) -> Self: ...
+                    def example_list(self) -> List[Self]: ...
+                NewManager = BaseManager.from_queryset(ModelQuerySet)
+                class MyModel(models.Model):
+                    objects = NewManager()
+
 -   case: from_queryset_with_base_manager
     main: |
         from myapp.models import MyModel