From 719f7edd20a296ee64c6243b572be3ca96005b6b Mon Sep 17 00:00:00 2001
From: moran abadie <moran.abadie@gmail.com>
Date: Wed, 30 Aug 2023 18:31:50 +0200
Subject: [PATCH] Add TypeVar resolution for as_manager and from_queryset
 querysets (typeddjango#1646)

---
 mypy_django_plugin/transformers/managers.py   | 21 ++++++++++-
 .../managers/querysets/test_as_manager.yml    | 26 ++++++++++++++
 .../managers/querysets/test_from_queryset.yml | 35 +++++++++++++++++++
 3 files changed, 81 insertions(+), 1 deletion(-)

diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py
index 7fa3244373..f9386be225 100644
--- a/mypy_django_plugin/transformers/managers.py
+++ b/mypy_django_plugin/transformers/managers.py
@@ -18,7 +18,7 @@
 from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext
 from mypy.semanal import SemanticAnalyzer
 from mypy.semanal_shared import has_placeholder
-from mypy.types import AnyType, CallableType, Instance, ProperType, TypeOfAny
+from mypy.types import AnyType, CallableType, Instance, ProperType, TypeOfAny, TypeVarType
 from mypy.types import Type as MypyType
 from mypy.typevars import fill_typevars
 
@@ -105,6 +105,9 @@ def get_funcdef_type(definition: Union[FuncBase, Decorator, None]) -> Optional[P
             ret_type = Instance(queryset_info, [manager_instance.args[0], manager_instance.args[0]])
         variables = []
 
+    if isinstance(ret_type, TypeVarType):
+        ret_type = _find_type_var(queryset_info, ret_type) or ret_type
+
     # Drop any 'self' argument as our manager is already initialized
     return method_type.copy_modified(
         arg_types=method_type.arg_types[1:],
@@ -115,6 +118,22 @@ def get_funcdef_type(definition: Union[FuncBase, Decorator, None]) -> Optional[P
     )
 
 
+def _find_type_var(queryset_info: TypeInfo, ret_type: TypeVarType) -> Optional[MypyType]:
+    """
+    Attempts to find the concrete type corresponding to a TypeVarType in a queryset's base types.
+
+    Example:
+    Suppose queryset_info is based on a generic class `QuerySet[T]` and ret_type corresponds
+    to `T`, then this function will return the concrete type that `T` was instantiated with.
+    """
+    for base in queryset_info.bases:
+        for i, type_var in enumerate(base.type.type_vars):
+            type_var_path = f"{base.type.module_name}.{type_var}"
+            if type_var_path == ret_type.fullname and i < len(base.args):
+                return base.args[i]
+    return None
+
+
 def get_method_type_from_reverse_manager(
     api: TypeChecker, method_name: str, manager_type_info: TypeInfo
 ) -> Optional[ProperType]:
diff --git a/tests/typecheck/managers/querysets/test_as_manager.yml b/tests/typecheck/managers/querysets/test_as_manager.yml
index 06882368b9..e49ffd84b1 100644
--- a/tests/typecheck/managers/querysets/test_as_manager.yml
+++ b/tests/typecheck/managers/querysets/test_as_manager.yml
@@ -146,6 +146,32 @@
                 class MyModel(models.Model):
                     objects = ManagerFromModelQuerySet
 
+-   case: handles_subclasses_of_queryset
+    main: |
+        from myapp.models import MyModel
+        reveal_type(MyModel.objects.example())  # N: Revealed type is "myapp.models.MyModel"
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from typing import TypeVar
+
+                from django.db import models
+
+                _CTE = TypeVar("_CTE", bound=models.Model)
+
+                class _MyModelQuerySet(models.QuerySet[_CTE]):
+
+                    def example(self) -> _CTE: ...
+
+                class MyModelQuerySet(_MyModelQuerySet["MyModel"]):
+                    ...
+
+                class MyModel(models.Model):
+                    objects = MyModelQuerySet.as_manager()
+
 -   case: reuses_generated_type_when_called_identically_for_multiple_managers
     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 aa152983f5..f9c5af288f 100644
--- a/tests/typecheck/managers/querysets/test_from_queryset.yml
+++ b/tests/typecheck/managers/querysets/test_from_queryset.yml
@@ -70,6 +70,41 @@
                 NewManager = BaseManager.from_queryset(ModelQuerySet)
                 class MyModel(models.Model):
                     objects = NewManager()
+-   case: handles_subclasses_of_queryset
+    main: |
+        from myapp.models import MyModel
+        reveal_type(MyModel.objects.example())  # N: Revealed type is "myapp.models.MyModel"
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/querysets.py
+            content: |
+                from typing import TypeVar, TYPE_CHECKING
+
+                from django.db import models
+                from django.db.models.manager import BaseManager
+                if TYPE_CHECKING:
+                    from .models import MyModel
+
+                _CTE = TypeVar("_CTE", bound=models.Model)
+
+                class _MyModelQuerySet(models.QuerySet[_CTE]):
+
+                    def example(self) -> _CTE: ...
+
+                class MyModelQuerySet(_MyModelQuerySet["MyModel"]):
+                    ...
+
+        - path: myapp/models.py
+          content: |
+            from django.db import models
+            from django.db.models.manager import BaseManager
+            from .querysets import MyModelQuerySet
+
+            NewManager = BaseManager.from_queryset(MyModelQuerySet)
+            class MyModel(models.Model):
+                objects = NewManager()
 
 -   case: from_queryset_generated_manager_imported_from_other_module
     main: |