Fix IndexError crash when using from_queryset() of custom Manager subclass (#1786)

Author: moranabadie

mypy_django_plugin/transformers/managers.py

@@ -125,13 +125,17 @@ def _process_dynamic_method(
         variables = []
     args_types = method_type.arg_types[1:]
     if _has_compatible_type_vars(base_that_has_method):
-        ret_type = _replace_type_var(
-            ret_type, base_that_has_method.defn.type_vars[0].fullname, manager_instance.args[0]
-        )
-        args_types = [
-            _replace_type_var(arg_type, base_that_has_method.defn.type_vars[0].fullname, manager_instance.args[0])
-            for arg_type in args_types
-        ]
+        typed_var = manager_instance.args or queryset_info.bases[0].args
+        if (
+            typed_var
+            and isinstance(typed_var[0], Instance)
+            and typed_var[0].type.has_base(fullnames.MODEL_CLASS_FULLNAME)
+        ):
+            ret_type = _replace_type_var(ret_type, base_that_has_method.defn.type_vars[0].fullname, typed_var[0])
+            args_types = [
+                _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)

tests/typecheck/managers/querysets/test_from_queryset.yml

@@ -93,6 +93,43 @@
                 NewManager = BaseManager.from_queryset(ModelQuerySet)
                 class MyModel(models.Model):
                     objects = NewManager()
+
+-   case: from_queryset_custom_manager_subclass
+    main: |
+        from myapp.models import NewManager
+        reveal_type(NewManager().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
+
+                class TypedManager(BaseManager["MyModel"]): ...
+
+                NewManager = TypedManager.from_queryset(MyModelQuerySet)
+                class MyModel(models.Model):
+                    objects = NewManager()
 -   case: handles_subclasses_of_queryset
     main: |
         from myapp.models import MyModel