From 04c37a70bed895e8443fa0f92a6fba8660d0350b Mon Sep 17 00:00:00 2001
From: moran abadie <moran.abadie@gmail.com>
Date: Fri, 20 Oct 2023 17:43:03 +0200
Subject: [PATCH] fix as manager and from queryset self types (#1788)

---
 mypy_django_plugin/transformers/managers.py   |  3 ++
 .../managers/querysets/test_as_manager.yml    | 28 +++++++++++++++++++
 .../managers/querysets/test_from_queryset.yml | 23 +++++++++++++++
 3 files changed, 54 insertions(+)

diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py
index 4d362c1d3..c4b67e6fa 100644
--- 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
index 7d6e92271..9ef966aca 100644
--- 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
index ec08da3da..8c1701f5f 100644
--- 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