Improve ManyToManyDescriptor and fix Model.<manytomany>.through typing (#1805)

`ManyToManyDescriptor` is now extended to take 1 new type argument, which is the target model/other side of the relation.

The plugin is updated to:

- Set a `ManyToManyDescriptor` instance instead of a related manager, for reverse relations of a `ManyToManyField`
- Produce a manager class with `ManyRelatedManager` and a model's default manager as bases for both sides of a many-to-many relation

Author: flaeppe

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

@@ -25,7 +25,6 @@ RECURSIVE_RELATIONSHIP_CONSTANT: Literal["self"]
 
 def resolve_relation(scope_model: type[Model], relation: str | type[Model]) -> str | type[Model]: ...
 
-_M = TypeVar("_M", bound=Model)
 # __set__ value type
 _ST = TypeVar("_ST")
 # __get__ return type
@@ -232,9 +231,10 @@ class OneToOneField(ForeignKey[_ST, _GT]):
     @overload
     def __get__(self, instance: Any, owner: Any) -> Self: ...
 
+_Through = TypeVar("_Through", bound=Model)
 _To = TypeVar("_To", bound=Model)
 
-class ManyToManyField(RelatedField[Any, Any], Generic[_To, _M]):
+class ManyToManyField(RelatedField[Any, Any], Generic[_To, _Through]):
     description: str
     has_null_arg: bool
     swappable: bool
@@ -253,7 +253,7 @@ class ManyToManyField(RelatedField[Any, Any], Generic[_To, _M]):
         related_query_name: str | None = ...,
         limit_choices_to: _AllLimitChoicesTo | None = ...,
         symmetrical: bool | None = ...,
-        through: type[_M] | str | None = ...,
+        through: type[_Through] | str | None = ...,
         through_fields: tuple[str, str] | None = ...,
         db_constraint: bool = ...,
         db_table: str | None = ...,
@@ -282,7 +282,7 @@ class ManyToManyField(RelatedField[Any, Any], Generic[_To, _M]):
     ) -> None: ...
     # class access
     @overload
-    def __get__(self, instance: None, owner: Any) -> ManyToManyDescriptor[_M]: ...
+    def __get__(self, instance: None, owner: Any) -> ManyToManyDescriptor[_To, _Through]: ...
     # Model instance access
     @overload
     def __get__(self, instance: Model, owner: Any) -> ManyRelatedManager[_To]: ...

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

@@ -15,6 +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)
 _To = TypeVar("_To", bound=Model)
 
 class ForeignKeyDeferredAttribute(DeferredAttribute):
@@ -84,7 +85,7 @@ class ReverseManyToOneDescriptor:
     @overload
     def __get__(self, instance: None, cls: Any = ...) -> Self: ...
     @overload
-    def __get__(self, instance: Model, cls: Any = ...) -> type[RelatedManager[Any]]: ...
+    def __get__(self, instance: Model, cls: Any = ...) -> RelatedManager[Any]: ...
     def __set__(self, instance: Any, value: Any) -> NoReturn: ...
 
 # Fake class, Django defines 'RelatedManager' inside a function body
@@ -104,7 +105,7 @@ def create_reverse_many_to_one_manager(
     superclass: type[BaseManager[_M]], rel: ManyToOneRel
 ) -> type[RelatedManager[_M]]: ...
 
-class ManyToManyDescriptor(ReverseManyToOneDescriptor, Generic[_M]):
+class ManyToManyDescriptor(ReverseManyToOneDescriptor, Generic[_To, _Through]):
     """
     In the example::
 
@@ -117,13 +118,17 @@ class ManyToManyDescriptor(ReverseManyToOneDescriptor, Generic[_M]):
 
     # 'field' here is 'rel.field'
     rel: ManyToManyRel  # type: ignore[assignment]
-    field: ManyToManyField[Any, _M]  # type: ignore[assignment]
+    field: ManyToManyField[_To, _Through]  # type: ignore[assignment]
     reverse: bool
     def __init__(self, rel: ManyToManyRel, reverse: bool = ...) -> None: ...
     @property
-    def through(self) -> type[_M]: ...
+    def through(self) -> type[_Through]: ...
     @cached_property
-    def related_manager_cls(self) -> type[ManyRelatedManager[Any]]: ...  # type: ignore[override]
+    def related_manager_cls(self) -> type[ManyRelatedManager[_To]]: ...  # 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]: ...
 
 # Fake class, Django defines 'ManyRelatedManager' inside a function body
 class ManyRelatedManager(Manager[_M], Generic[_M]):

mypy_django_plugin/lib/fullnames.py

@@ -34,6 +34,8 @@
 }
 
 REVERSE_ONE_TO_ONE_DESCRIPTOR = "django.db.models.fields.related_descriptors.ReverseOneToOneDescriptor"
+MANY_TO_MANY_DESCRIPTOR = "django.db.models.fields.related_descriptors.ManyToManyDescriptor"
+MANY_RELATED_MANAGER = "django.db.models.fields.related_descriptors.ManyRelatedManager"
 RELATED_FIELDS_CLASSES = frozenset(
     (
         FOREIGN_OBJECT_FULLNAME,

mypy_django_plugin/lib/helpers.py

@@ -67,6 +67,20 @@ def get_django_metadata_bases(
     return get_django_metadata(model_info).setdefault(key, cast(Dict[str, int], {}))
 
 
+def get_reverse_manager_info(
+    api: Union[TypeChecker, SemanticAnalyzer], model_info: TypeInfo, derived_from: str
+) -> Optional[TypeInfo]:
+    manager_fullname = get_django_metadata(model_info).get("reverse_managers", {}).get(derived_from)
+    if not manager_fullname:
+        return None
+
+    return lookup_fully_qualified_typeinfo(api, manager_fullname)
+
+
+def set_reverse_manager_info(model_info: TypeInfo, derived_from: str, fullname: str) -> None:
+    get_django_metadata(model_info).setdefault("reverse_managers", {})[derived_from] = fullname
+
+
 class IncompleteDefnException(Exception):
     pass
 
@@ -457,3 +471,10 @@ def resolve_lazy_reference(
     else:
         api.fail("Could not match lazy reference with any model", ctx)
     return None
+
+
+def is_model_instance(instance: Instance) -> bool:
+    return (
+        instance.type.metaclass_type is not None
+        and instance.type.metaclass_type.type.fullname == fullnames.MODEL_METACLASS_FULLNAME
+    )

mypy_django_plugin/main.py

@@ -23,7 +23,7 @@
 from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.exceptions import UnregisteredModelError
 from mypy_django_plugin.lib import fullnames, helpers
-from mypy_django_plugin.transformers import fields, forms, init_create, meta, querysets, request, settings
+from mypy_django_plugin.transformers import fields, forms, init_create, manytomany, meta, querysets, request, settings
 from mypy_django_plugin.transformers.functional import resolve_str_promise_attribute
 from mypy_django_plugin.transformers.managers import (
     create_new_manager_class_from_as_manager_method,
@@ -188,6 +188,12 @@ def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], M
             if info and info.has_base(fullnames.FORM_MIXIN_CLASS_FULLNAME):
                 return forms.extract_proper_type_for_get_form
 
+        elif method_name == "__get__" and class_fullname in {
+            fullnames.MANYTOMANY_FIELD_FULLNAME,
+            fullnames.MANY_TO_MANY_DESCRIPTOR,
+        }:
+            return manytomany.refine_many_to_many_related_manager
+
         manager_classes = self._get_current_manager_bases()
 
         if method_name == "values":

mypy_django_plugin/transformers/manytomany.py

@@ -1,8 +1,8 @@
-from typing import NamedTuple, Optional, Union
+from typing import NamedTuple, Optional, Tuple, Union
 
 from mypy.checker import TypeChecker
 from mypy.nodes import AssignmentStmt, Expression, MemberExpr, NameExpr, StrExpr, TypeInfo
-from mypy.plugin import FunctionContext
+from mypy.plugin import FunctionContext, MethodContext
 from mypy.semanal import SemanticAnalyzer
 from mypy.types import Instance, ProperType, UninhabitedType
 from mypy.types import Type as MypyType
@@ -151,3 +151,57 @@ def get_model_from_expression(
         if model_info is not None:
             return Instance(model_info, [])
     return None
+
+
+def get_related_manager_and_model(ctx: MethodContext) -> Optional[Tuple[Instance, Instance]]:
+    if (
+        isinstance(ctx.default_return_type, Instance)
+        and ctx.default_return_type.type.fullname == fullnames.MANY_RELATED_MANAGER
+    ):
+        # This is a call to '__get__' overload with a model instance of 'ManyToManyDescriptor'.
+        # 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
+        if (
+            many_related_manager.args
+            and isinstance(many_related_manager.args[0], Instance)
+            and helpers.is_model_instance(many_related_manager.args[0])
+        ):
+            return many_related_manager, many_related_manager.args[0]
+
+    return None
+
+
+def refine_many_to_many_related_manager(ctx: MethodContext) -> MypyType:
+    """
+    Updates the 'ManyRelatedManager' returned by e.g. 'ManyToManyDescriptor' to be a subclass
+    of 'ManyRelatedManager' and the related model's default manager.
+    """
+    related_objects = get_related_manager_and_model(ctx)
+    if related_objects is None:
+        return ctx.default_return_type
+
+    many_related_manager, related_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"
+    )
+    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
+
+        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],
+        )
+        related_manager_info.metadata["django"] = {"related_manager_to_model": related_model_instance.type.fullname}
+        helpers.set_reverse_manager_info(
+            related_model_instance.type,
+            derived_from="_default_manager",
+            fullname=related_manager_info.fullname,
+        )
+    return Instance(related_manager_info, [])

mypy_django_plugin/transformers/models.py

@@ -3,7 +3,7 @@
 
 from django.db.models import Manager, Model
 from django.db.models.fields import DateField, DateTimeField, Field
-from django.db.models.fields.reverse_related import ForeignObjectRel, OneToOneRel
+from django.db.models.fields.reverse_related import ManyToManyRel, OneToOneRel
 from mypy.checker import TypeChecker
 from mypy.nodes import (
     ARG_STAR2,
@@ -448,23 +448,15 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
 
 
 class AddReverseLookups(ModelClassInitializer):
-    def get_reverse_manager_info(self, model_info: TypeInfo, derived_from: str) -> Optional[TypeInfo]:
-        manager_fullname = helpers.get_django_metadata(model_info).get("reverse_managers", {}).get(derived_from)
-        if not manager_fullname:
-            return None
-
-        symbol = self.api.lookup_fully_qualified_or_none(manager_fullname)
-        if symbol is None or not isinstance(symbol.node, TypeInfo):
-            return None
-        return symbol.node
+    @cached_property
+    def reverse_one_to_one_descriptor(self) -> TypeInfo:
+        return self.lookup_typeinfo_or_incomplete_defn_error(fullnames.REVERSE_ONE_TO_ONE_DESCRIPTOR)
 
-    def set_reverse_manager_info(self, model_info: TypeInfo, derived_from: str, fullname: str) -> None:
-        helpers.get_django_metadata(model_info).setdefault("reverse_managers", {})[derived_from] = fullname
+    @cached_property
+    def many_to_many_descriptor(self) -> TypeInfo:
+        return self.lookup_typeinfo_or_incomplete_defn_error(fullnames.MANY_TO_MANY_DESCRIPTOR)
 
     def run_with_model_cls(self, model_cls: Type[Model]) -> None:
-        reverse_one_to_one_descriptor = self.lookup_typeinfo_or_incomplete_defn_error(
-            fullnames.REVERSE_ONE_TO_ONE_DESCRIPTOR
-        )
         # add related managers
         for relation in self.django_context.get_model_relations(model_cls):
             attname = relation.get_accessor_name()
@@ -487,13 +479,27 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
                 self.add_new_node_to_model_class(
                     attname,
                     Instance(
-                        reverse_one_to_one_descriptor,
+                        self.reverse_one_to_one_descriptor,
                         [Instance(self.model_classdef.info, []), Instance(related_model_info, [])],
                     ),
                 )
                 continue
 
-            if isinstance(relation, ForeignObjectRel):
+            elif isinstance(relation, ManyToManyRel):
+                # TODO: 'relation' should be based on `TypeInfo` instead of Django runtime.
+                to_fullname = helpers.get_class_fullname(relation.remote_field.model)
+                to_model_info = self.lookup_typeinfo_or_incomplete_defn_error(to_fullname)
+                assert relation.through is not None
+                through_fullname = helpers.get_class_fullname(relation.through)
+                through_model_info = self.lookup_typeinfo_or_incomplete_defn_error(through_fullname)
+                self.add_new_node_to_model_class(
+                    attname,
+                    Instance(
+                        self.many_to_many_descriptor, [Instance(to_model_info, []), Instance(through_model_info, [])]
+                    ),
+                )
+
+            else:
                 related_manager_info = None
                 try:
                     related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(
@@ -534,8 +540,8 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
 
                 # Check if the related model has a related manager subclassed from the default manager
                 # TODO: Support other reverse managers than `_default_manager`
-                default_reverse_manager_info = self.get_reverse_manager_info(
-                    model_info=related_model_info, derived_from="_default_manager"
+                default_reverse_manager_info = helpers.get_reverse_manager_info(
+                    self.api, model_info=related_model_info, derived_from="_default_manager"
                 )
                 if default_reverse_manager_info:
                     self.add_new_node_to_model_class(attname, Instance(default_reverse_manager_info, []))
@@ -564,7 +570,7 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
                 new_related_manager_info.metadata["django"] = {"related_manager_to_model": related_model_info.fullname}
                 # Stash the new reverse manager type fullname on the related model, so we don't duplicate
                 # or have to create it again for other reverse relations
-                self.set_reverse_manager_info(
+                helpers.set_reverse_manager_info(
                     related_model_info,
                     derived_from="_default_manager",
                     fullname=new_related_manager_info.fullname,

mypy_django_plugin/transformers/orm_lookups.py

@@ -15,9 +15,7 @@ def typecheck_queryset_filter(ctx: MethodContext, django_context: DjangoContext)
     lookup_kwargs = ctx.arg_names[1] if len(ctx.arg_names) >= 2 else []
     provided_lookup_types = ctx.arg_types[1] if len(ctx.arg_types) >= 2 else []
 
-    assert isinstance(ctx.type, Instance)
-
-    if not ctx.type.args or not isinstance(ctx.type.args[0], Instance):
+    if not isinstance(ctx.type, Instance) or not ctx.type.args or not isinstance(ctx.type.args[0], Instance):
         return ctx.default_return_type
 
     model_cls_fullname = ctx.type.args[0].type.fullname

scripts/stubtest/allowlist.txt

@@ -26,9 +26,7 @@ django.db.models.fields.related_descriptors.RelatedManager
 # _locally/dynamically_ runtime -- Created via
 # 'django.db.models.fields.related_descriptors.create_reverse_many_to_one_manager'
 django.contrib.admin.models.LogEntry_RelatedManager
-django.contrib.auth.models.Group_RelatedManager
 django.contrib.auth.models.Permission_RelatedManager
-django.contrib.auth.models.User_RelatedManager
 
 # BaseArchive abstract methods that take no argument, but typed with arguments to match the Archive and TarArchive Implementations
 django.utils.archive.BaseArchive.list

scripts/stubtest/allowlist_todo.txt

@@ -167,7 +167,6 @@ django.contrib.auth.models.GroupManager.__slotnames__
 django.contrib.auth.models.Permission.codename
 django.contrib.auth.models.Permission.content_type
 django.contrib.auth.models.Permission.content_type_id
-django.contrib.auth.models.Permission.group_set
 django.contrib.auth.models.Permission.id
 django.contrib.auth.models.Permission.name
 django.contrib.auth.models.Permission.user_set