Add TypeVar resolution for as_manager querysets (typeddjango#1646)

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]:

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