Improve hints for ManyToManyDescriptor and start using it

Author: flaeppe

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

@@ -253,7 +253,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, _M]: ...
     # Model instance access
     @overload
     def __get__(self, instance: Model, owner: Any) -> ManyRelatedManager[_To]: ...

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

@@ -14,6 +14,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):
@@ -83,14 +84,14 @@ 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: ...
 
 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::
 
@@ -103,13 +104,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]: ...
     @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]: ...
 
 class ManyRelatedManager(BaseManager[_M], Generic[_M]):
     related_val: tuple[int, ...]
@@ -121,7 +126,7 @@ class ManyRelatedManager(BaseManager[_M], Generic[_M]):
     async def aset(self, objs: QuerySet[_M] | Iterable[_M | int], *, bulk: bool = ..., clear: bool = ...) -> None: ...
     def clear(self) -> None: ...
     async def aclear(self) -> None: ...
-    def __call__(self, *, manager: str) -> ManyRelatedManager[_M]: ...
+    def __call__(self, *, manager: str) -> Self: ...
 
 def create_forward_many_to_many_manager(
     superclass: type[BaseManager[_M]], rel: ManyToManyRel, reverse: bool

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
 

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,
@@ -187,6 +187,11 @@ def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], M
             info = self._get_typeinfo_or_none(class_fullname)
             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()
 

mypy_django_plugin/transformers/manytomany.py

@@ -2,7 +2,7 @@
 
 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,52 @@ def get_model_from_expression(
         if model_info is not None:
             return Instance(model_info, [])
     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.
+    """
+    if (
+        not isinstance(ctx.default_return_type, Instance)
+        # Only change a return type being 'ManyRelatedManager'
+        or ctx.default_return_type.type.fullname != fullnames.MANY_RELATED_MANAGER
+    ):
+        return ctx.default_return_type
+
+    # 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 (
+        not many_related_manager.args
+        or not isinstance(many_related_manager.args[0], Instance)
+        or many_related_manager.args[0].type.metaclass_type is None
+        or many_related_manager.args[0].type.metaclass_type.type.fullname != fullnames.MODEL_METACLASS_FULLNAME
+    ):
+        return ctx.default_return_type
+
+    checker = helpers.get_typechecker_api(ctx)
+    related_model_instance = many_related_manager.args[0].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,25 @@ 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 +538,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 +568,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_todo.txt

@@ -170,7 +170,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