Add ManyRelatedManager.through attribute and generic type parameter

django-stubs/db/models/fields/related.pyi

@@ -296,7 +296,7 @@ class ManyToManyField(RelatedField[Any, Any], Generic[_To, _Through]):
     def __get__(self, instance: None, owner: Any) -> ManyToManyDescriptor[_To, _Through]: ...
     # Model instance access
     @overload
-    def __get__(self, instance: Model, owner: Any) -> ManyRelatedManager[_To]: ...
+    def __get__(self, instance: Model, owner: Any) -> ManyRelatedManager[_To, _Through]: ...
     # non-Model instances
     @overload
     def __get__(self, instance: Any, owner: Any) -> Self: ...

django-stubs/db/models/fields/related_descriptors.pyi

@@ -15,7 +15,7 @@ from typing_extensions import Self
 _M = TypeVar("_M", bound=Model)
 _F = TypeVar("_F", bound=Field)
 _From = TypeVar("_From", bound=Model)
-_Through = TypeVar("_Through", bound=Model)
+_Through = TypeVar("_Through", bound=Model, default=Model)
 _To = TypeVar("_To", bound=Model)
 
 class ForeignKeyDeferredAttribute(DeferredAttribute):
@@ -125,40 +125,41 @@ class ManyToManyDescriptor(ReverseManyToOneDescriptor, Generic[_To, _Through]):
     @property
     def through(self) -> type[_Through]: ...
     @cached_property
-    def related_manager_cls(self) -> type[ManyRelatedManager[_To]]: ...  # type: ignore[override]
+    def related_manager_cls(self) -> type[ManyRelatedManager[_To, _Through]]: ...  # type: ignore[override]
     @overload  # type: ignore[override]
     def __get__(self, instance: None, cls: Any = ...) -> Self: ...
     @overload
-    def __get__(self, instance: Model, cls: Any = ...) -> ManyRelatedManager[_To]: ...
+    def __get__(self, instance: Model, cls: Any = ...) -> ManyRelatedManager[_To, _Through]: ...
 
 # Fake class, Django defines 'ManyRelatedManager' inside a function body
 @type_check_only
-class ManyRelatedManager(Manager[_M], Generic[_M]):
+class ManyRelatedManager(Manager[_To], Generic[_To, _Through]):
     related_val: tuple[int, ...]
-    def add(self, *objs: _M | int, bulk: bool = ..., through_defaults: dict[str, Any] | None = ...) -> None: ...
-    async def aadd(self, *objs: _M | int, bulk: bool = ..., through_defaults: dict[str, Any] | None = ...) -> None: ...
-    def remove(self, *objs: _M | int, bulk: bool = ...) -> None: ...
-    async def aremove(self, *objs: _M | int, bulk: bool = ...) -> None: ...
+    through: type[_Through]
+    def add(self, *objs: _To | int, bulk: bool = ..., through_defaults: dict[str, Any] | None = ...) -> None: ...
+    async def aadd(self, *objs: _To | int, bulk: bool = ..., through_defaults: dict[str, Any] | None = ...) -> None: ...
+    def remove(self, *objs: _To | int, bulk: bool = ...) -> None: ...
+    async def aremove(self, *objs: _To | int, bulk: bool = ...) -> None: ...
     def set(
         self,
-        objs: QuerySet[_M] | Iterable[_M | int],
+        objs: QuerySet[_To] | Iterable[_To | int],
         *,
         bulk: bool = ...,
         clear: bool = ...,
         through_defaults: dict[str, Any] | None = ...,
     ) -> None: ...
     async def aset(
         self,
-        objs: QuerySet[_M] | Iterable[_M | int],
+        objs: QuerySet[_To] | Iterable[_To | int],
         *,
         bulk: bool = ...,
         clear: bool = ...,
         through_defaults: dict[str, Any] | None = ...,
     ) -> None: ...
     def clear(self) -> None: ...
     async def aclear(self) -> None: ...
-    def __call__(self, *, manager: str) -> ManyRelatedManager[_M]: ...
+    def __call__(self, *, manager: str) -> ManyRelatedManager[_To, _Through]: ...
 
 def create_forward_many_to_many_manager(
-    superclass: type[BaseManager[_M]], rel: ManyToManyRel, reverse: bool
-) -> type[ManyRelatedManager[_M]]: ...
+    superclass: type[BaseManager[_To]], rel: ManyToManyRel, reverse: bool
+) -> type[ManyRelatedManager[_To, _Through]]: ...

ext/django_stubs_ext/db/models/manager.py

@@ -13,11 +13,12 @@
     from typing import Protocol, TypeVar
 
     _T = TypeVar("_T")
+    _Through = TypeVar("_Through")
 
     # Define as `Protocol` to prevent them being used with `isinstance()`.
     # These actually inherit from `BaseManager`.
     class RelatedManager(Protocol[_T]):
         pass
 
-    class ManyRelatedManager(Protocol[_T]):
+    class ManyRelatedManager(Protocol[_T, _Through]):
         pass

mypy_django_plugin/lib/helpers.py

@@ -55,6 +55,7 @@ class DjangoTypeMetadata(TypedDict, total=False):
     model_bases: Dict[str, int]
     queryset_bases: Dict[str, int]
     m2m_throughs: Dict[str, str]
+    m2m_managers: Dict[str, str]
 
 
 def get_django_metadata(model_info: TypeInfo) -> DjangoTypeMetadata:
@@ -81,6 +82,20 @@ def set_reverse_manager_info(model_info: TypeInfo, derived_from: str, fullname:
     get_django_metadata(model_info).setdefault("reverse_managers", {})[derived_from] = fullname
 
 
+def get_many_to_many_manager_info(
+    api: Union[TypeChecker, SemanticAnalyzer], *, to: TypeInfo, derived_from: str
+) -> Optional[TypeInfo]:
+    manager_fullname = get_django_metadata(to).get("m2m_managers", {}).get(derived_from)
+    if not manager_fullname:
+        return None
+
+    return lookup_fully_qualified_typeinfo(api, manager_fullname)
+
+
+def set_many_to_many_manager_info(to: TypeInfo, derived_from: str, manager_info: TypeInfo) -> None:
+    get_django_metadata(to).setdefault("m2m_managers", {})[derived_from] = manager_info.fullname
+
+
 class IncompleteDefnException(Exception):
     pass
 

mypy_django_plugin/transformers/manytomany.py

@@ -4,7 +4,7 @@
 from mypy.nodes import AssignmentStmt, Expression, MemberExpr, NameExpr, StrExpr, TypeInfo
 from mypy.plugin import FunctionContext, MethodContext
 from mypy.semanal import SemanticAnalyzer
-from mypy.types import Instance, ProperType, UninhabitedType
+from mypy.types import Instance, ProperType, TypeVarType, UninhabitedType
 from mypy.types import Type as MypyType
 
 from mypy_django_plugin.django.context import DjangoContext
@@ -150,7 +150,18 @@ def get_model_from_expression(
     return None
 
 
-def get_related_manager_and_model(ctx: MethodContext) -> Optional[Tuple[Instance, Instance]]:
+def get_related_manager_and_model(ctx: MethodContext) -> Optional[Tuple[Instance, Instance, Instance]]:
+    """
+    Returns a 3-tuple consisting of:
+      1. A `ManyRelatedManager` instance
+      2. The first type parameter (_To) instance of 1. when it's a model
+      3. The second type parameter (_Through) instance of 1. when it's a model
+    When encountering a `ManyRelatedManager` that has populated its 2 first type
+    parameters with models. Otherwise `None` is returned.
+
+    For example: if given a `ManyRelatedManager[A, B]` where `A` and `B` are models the
+    following 3-tuple is returned: `(ManyRelatedManager[A, B], A, B)`.
+    """
     if (
         isinstance(ctx.default_return_type, Instance)
         and ctx.default_return_type.type.fullname == fullnames.MANY_RELATED_MANAGER
@@ -159,13 +170,15 @@ def get_related_manager_and_model(ctx: MethodContext) -> Optional[Tuple[Instance
         # Returning a 'ManyRelatedManager'. Which we want to, just like Django, build from the
         # default manager of the related model.
         many_related_manager = ctx.default_return_type
-        # Require first type argument of 'ManyRelatedManager' to be a model
+        # Require first and second type argument of 'ManyRelatedManager' to be models
         if (
-            many_related_manager.args
+            len(many_related_manager.args) >= 2
             and isinstance(many_related_manager.args[0], Instance)
             and helpers.is_model_type(many_related_manager.args[0].type)
+            and isinstance(many_related_manager.args[1], Instance)
+            and helpers.is_model_type(many_related_manager.args[1].type)
         ):
-            return many_related_manager, many_related_manager.args[0]
+            return many_related_manager, many_related_manager.args[0], many_related_manager.args[1]
 
     return None
 
@@ -179,26 +192,46 @@ def refine_many_to_many_related_manager(ctx: MethodContext) -> MypyType:
     if related_objects is None:
         return ctx.default_return_type
 
-    many_related_manager, related_model_instance = related_objects
+    many_related_manager, related_model_instance, through_model_instance = related_objects
     checker = helpers.get_typechecker_api(ctx)
-    related_model_instance = related_model_instance.copy_modified()
-    related_manager_info = helpers.get_reverse_manager_info(
-        checker, related_model_instance.type, derived_from="_default_manager"
+    related_manager_info = helpers.get_many_to_many_manager_info(
+        checker, to=related_model_instance.type, derived_from="_default_manager"
     )
     if related_manager_info is None:
         default_manager_node = related_model_instance.type.names.get("_default_manager")
         if default_manager_node is None or not isinstance(default_manager_node.type, Instance):
             return ctx.default_return_type
 
+        # Create a reusable generic subclass that is generic over a 'through' model,
+        # explicitly declared it'd could have looked something like below
+        #
+        # class X(models.Model): ...
+        # _Through = TypeVar("_Through", bound=models.Model)
+        # class X_ManyRelatedManager(ManyRelatedManager[X, _Through], type(X._default_manager), Generic[_Through]): ...
+        _through_type_var = many_related_manager.type.defn.type_vars[1]
+        assert isinstance(_through_type_var, TypeVarType)
+        generic_to_many_related_manager = many_related_manager.copy_modified(
+            args=[
+                # Keep the same '_To' as the (parent) `ManyRelatedManager` instance
+                many_related_manager.args[0],
+                # But reset the '_Through' `TypeVar` declared for `ManyRelatedManager`
+                _through_type_var.copy_modified(),
+            ]
+        )
         related_manager_info = helpers.add_new_class_for_module(
             module=checker.modules[related_model_instance.type.module_name],
             name=f"{related_model_instance.type.name}_ManyRelatedManager",
-            bases=[many_related_manager, default_manager_node.type],
+            bases=[generic_to_many_related_manager, default_manager_node.type],
         )
+        # Reuse the '_Through' `TypeVar` from `ManyRelatedManager` in our subclass
+        related_manager_info.defn.type_vars = [_through_type_var.copy_modified()]
+        related_manager_info.add_type_vars()
         related_manager_info.metadata["django"] = {"related_manager_to_model": related_model_instance.type.fullname}
-        helpers.set_reverse_manager_info(
-            related_model_instance.type,
+        # Track the existence of our manager subclass, by tying it to model it operates on
+        helpers.set_many_to_many_manager_info(
+            to=related_model_instance.type,
             derived_from="_default_manager",
-            fullname=related_manager_info.fullname,
+            manager_info=related_manager_info,
         )
-    return Instance(related_manager_info, [])
+
+    return Instance(related_manager_info, [through_model_instance])