From da86f3afec8be5f48f3625e44d806fa891caa631 Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Mon, 30 Oct 2023 09:12:40 +0100 Subject: [PATCH 1/7] Improve hints for `ManyToManyDescriptor` and start using it --- django-stubs/db/models/fields/related.pyi | 2 +- .../db/models/fields/related_descriptors.pyi | 17 ++++-- mypy_django_plugin/lib/fullnames.py | 2 + mypy_django_plugin/lib/helpers.py | 14 +++++ mypy_django_plugin/main.py | 7 ++- mypy_django_plugin/transformers/manytomany.py | 51 +++++++++++++++- mypy_django_plugin/transformers/models.py | 46 ++++++++------- .../transformers/orm_lookups.py | 4 +- scripts/stubtest/allowlist.txt | 2 - scripts/stubtest/allowlist_todo.txt | 1 - tests/typecheck/fields/test_related.yml | 58 ++++++++++++++----- .../typecheck/models/test_contrib_models.yml | 6 +- 12 files changed, 158 insertions(+), 52 deletions(-) diff --git a/django-stubs/db/models/fields/related.pyi b/django-stubs/db/models/fields/related.pyi index d7fdbd9d1..67525ac7e 100644 --- a/django-stubs/db/models/fields/related.pyi +++ b/django-stubs/db/models/fields/related.pyi @@ -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, _M]: ... # Model instance access @overload def __get__(self, instance: Model, owner: Any) -> ManyRelatedManager[_To]: ... diff --git a/django-stubs/db/models/fields/related_descriptors.pyi b/django-stubs/db/models/fields/related_descriptors.pyi index 75bf0caef..bb3f60d15 100644 --- a/django-stubs/db/models/fields/related_descriptors.pyi +++ b/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(BaseManager[_M], Generic[_M]): @@ -136,7 +141,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 diff --git a/mypy_django_plugin/lib/fullnames.py b/mypy_django_plugin/lib/fullnames.py index 33e746236..45fa3a452 100644 --- a/mypy_django_plugin/lib/fullnames.py +++ b/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, diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py index b262ea585..4d1da8c6c 100644 --- a/mypy_django_plugin/lib/helpers.py +++ b/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 diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py index 454a6cb13..19ed7c743 100644 --- a/mypy_django_plugin/main.py +++ b/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() diff --git a/mypy_django_plugin/transformers/manytomany.py b/mypy_django_plugin/transformers/manytomany.py index a0c812474..1f3b4715c 100644 --- a/mypy_django_plugin/transformers/manytomany.py +++ b/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, []) diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index 903d2b137..4920e67be 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/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, diff --git a/mypy_django_plugin/transformers/orm_lookups.py b/mypy_django_plugin/transformers/orm_lookups.py index 1bdf47d5f..7eed520fd 100644 --- a/mypy_django_plugin/transformers/orm_lookups.py +++ b/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 diff --git a/scripts/stubtest/allowlist.txt b/scripts/stubtest/allowlist.txt index 40376cb50..4f51fdfa9 100644 --- a/scripts/stubtest/allowlist.txt +++ b/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 diff --git a/scripts/stubtest/allowlist_todo.txt b/scripts/stubtest/allowlist_todo.txt index c9fb24a01..6ea4203fd 100644 --- a/scripts/stubtest/allowlist_todo.txt +++ b/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 diff --git a/tests/typecheck/fields/test_related.yml b/tests/typecheck/fields/test_related.yml index 36543b6c7..a0c4034b4 100644 --- a/tests/typecheck/fields/test_related.yml +++ b/tests/typecheck/fields/test_related.yml @@ -278,13 +278,16 @@ - case: many_to_many_field_converts_to_queryset_of_model_type main: | from myapp.models import App, Member - reveal_type(Member().apps) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyRelatedManager[myapp.models.App]" + reveal_type(Member().apps) # N: Revealed type is "myapp.models.App_ManyRelatedManager" reveal_type(Member().apps.get()) # N: Revealed type is "myapp.models.App" - reveal_type(App().members) # N: Revealed type is "django.db.models.fields.related_descriptors.RelatedManager[myapp.models.Member]" + reveal_type(Member().apps.app_method()) # N: Revealed type is "builtins.str" + Member().apps.add(App()) + reveal_type(App().members) # N: Revealed type is "myapp.models.Member_ManyRelatedManager" reveal_type(App().members.get()) # N: Revealed type is "myapp.models.Member" - reveal_type(Member.apps.field) # N: Revealed type is "django.db.models.fields.related.ManyToManyField[Any, myapp.models.Member_apps]" - # XXX the following is not correct: - reveal_type(App.members) # N: Revealed type is "django.db.models.fields.related_descriptors.RelatedManager[myapp.models.Member]" + reveal_type(App().members.member_method()) # N: Revealed type is "builtins.int" + App().members.add(Member()) + reveal_type(Member.apps.field) # N: Revealed type is "django.db.models.fields.related.ManyToManyField[myapp.models.App, myapp.models.Member_apps]" + reveal_type(App.members) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyToManyDescriptor[myapp.models.Member, myapp.models.Member_apps]" installed_apps: - myapp files: @@ -292,15 +295,23 @@ - path: myapp/models.py content: | from django.db import models + class AppQuerySet(models.QuerySet["App"]): + def app_method(self) -> str: + return "abc" class App(models.Model): - pass + objects = AppQuerySet.as_manager() + + class MemberQuerySet(models.QuerySet["Member"]): + def member_method(self) -> int: + return 1 class Member(models.Model): apps = models.ManyToManyField(to=App, related_name='members') + objects = MemberQuerySet.as_manager() - case: many_to_many_works_with_string_if_imported main: | from myapp.models import Member - reveal_type(Member().apps) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyRelatedManager[myapp2.models.App]" + reveal_type(Member().apps) # N: Revealed type is "myapp2.models.App_ManyRelatedManager" installed_apps: - myapp - myapp2 @@ -335,7 +346,8 @@ - case: many_to_many_with_self main: | from myapp.models import User - reveal_type(User().friends) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyRelatedManager[myapp.models.User]" + reveal_type(User().friends) # N: Revealed type is "myapp.models.User_ManyRelatedManager" + reveal_type(User().friends.get()) # N: Revealed type is "myapp.models.User" reveal_type(User.friends.through.objects.get()) # N: Revealed type is "myapp.models.User_friends" reveal_type(User.friends.through._default_manager) # N: Revealed type is "django.db.models.manager.Manager[myapp.models.User_friends]" reveal_type(User.friends.through().from_user) # N: Revealed type is "myapp.models.User" @@ -588,9 +600,9 @@ - case: test_related_fields_returned_as_descriptors_from_model_class main: | from myapp.models import Author, Blog, Publisher, Profile - reveal_type(Author.blogs) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyToManyDescriptor[myapp.models.Author_blogs]" + reveal_type(Author.blogs) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyToManyDescriptor[myapp.models.Blog, myapp.models.Author_blogs]" reveal_type(Author.blogs.through) # N: Revealed type is "Type[myapp.models.Author_blogs]" - reveal_type(Author().blogs) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyRelatedManager[myapp.models.Blog]" + reveal_type(Author().blogs) # N: Revealed type is "myapp.models.Blog_ManyRelatedManager" reveal_type(Blog.publisher) # N: Revealed type is "django.db.models.fields.related_descriptors.ForwardManyToOneDescriptor[django.db.models.fields.related.ForeignKey[Union[myapp.models.Publisher, django.db.models.expressions.Combinable], myapp.models.Publisher]]" reveal_type(Publisher.profile) # N: Revealed type is "django.db.models.fields.related_descriptors.ForwardOneToOneDescriptor[django.db.models.fields.related.OneToOneField[Union[myapp.models.Profile, django.db.models.expressions.Combinable], myapp.models.Profile]]" reveal_type(Author.file) # N: Revealed type is "django.db.models.fields.files.FileDescriptor" @@ -933,8 +945,8 @@ main: | from myapp.models import SalesMan sales_man = SalesMan() - reveal_type(sales_man.client) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyRelatedManager[myapp.models.CustomUser]" - reveal_type(sales_man.client(manager="staffs")) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyRelatedManager[myapp.models.CustomUser]" + reveal_type(sales_man.client) # N: Revealed type is "myapp.models.CustomUser_ManyRelatedManager" + reveal_type(sales_man.client(manager="staffs")) # N: Revealed type is "myapp.models.CustomUser_ManyRelatedManager" installed_apps: - myapp files: @@ -1092,6 +1104,13 @@ reveal_type(MyModel.auto_through.through._default_manager) reveal_type(MyModel.other_again.through) + reveal_type(Other.autos) + reveal_type(Other.autos.through) + reveal_type(Other.autos.through.objects.get()) + + reveal_type(Other().autos.custom_qs_method()) + reveal_type(Other().customs.custom_qs_method()) + reveal_type(Other().autos.add) out: | main:2: note: Revealed type is "myapp.models.MyModel_auto_through" main:3: note: Revealed type is "myapp.models.Other" @@ -1110,6 +1129,12 @@ main:20: note: Revealed type is "django.db.models.fields.related_descriptors.ForwardManyToOneDescriptor[django.db.models.fields.related.ForeignKey[Union[myapp.models.MyModel, django.db.models.expressions.Combinable], myapp.models.MyModel]]" main:21: note: Revealed type is "django.db.models.manager.Manager[myapp.models.MyModel_auto_through]" main:23: note: Revealed type is "Type[myapp.models.MyModel_other_again]" + main:24: note: Revealed type is "django.db.models.fields.related_descriptors.ManyToManyDescriptor[myapp.models.MyModel, myapp.models.MyModel_auto_through]" + main:25: note: Revealed type is "Type[myapp.models.MyModel_auto_through]" + main:26: note: Revealed type is "myapp.models.MyModel_auto_through" + main:28: note: Revealed type is "builtins.str" + main:29: note: Revealed type is "builtins.str" + main:30: note: Revealed type is "def (*objs: Union[myapp.models.MyModel, builtins.int], *, bulk: builtins.bool =)" installed_apps: - myapp files: @@ -1132,12 +1157,19 @@ objects = CustomThroughManager() + class MyModelQuerySet(models.QuerySet["MyModel"]): + def custom_qs_method(self) -> str: + return "abc" + + MyModelManager = models.Manager.from_queryset(MyModelQuerySet) class MyModel(models.Model): auto_through = models.ManyToManyField(Other, related_name="autos") # Have multiple M2Ms with implicit through other_again = models.ManyToManyField(Other, related_name="others_again") custom_through = models.ManyToManyField(Other, through=CustomThrough, related_name="customs") + objects = MyModelManager() + - case: test_many_to_many_with_lazy_references main: | from first.models import First @@ -1193,7 +1225,7 @@ reveal_type(Child.parents) reveal_type(Child().parents) out: | - main:2: note: Revealed type is "django.db.models.fields.related_descriptors.ManyToManyDescriptor[Any]" + main:2: note: Revealed type is "django.db.models.fields.related_descriptors.ManyToManyDescriptor[Any, Any]" main:3: note: Revealed type is "django.db.models.fields.related_descriptors.ManyRelatedManager[Any]" myapp/models/child:5: error: Need type annotation for "parents" [var-annotated] myapp/models/child:6: error: Need type annotation for "other_parents" [var-annotated] diff --git a/tests/typecheck/models/test_contrib_models.yml b/tests/typecheck/models/test_contrib_models.yml index 1b12124da..3daf86997 100644 --- a/tests/typecheck/models/test_contrib_models.yml +++ b/tests/typecheck/models/test_contrib_models.yml @@ -15,8 +15,8 @@ reveal_type(User().is_anonymous) # N: Revealed type is "Literal[False]" reveal_type(User().groups.get()) # N: Revealed type is "django.contrib.auth.models.Group" reveal_type(User().user_permissions.get()) # N: Revealed type is "django.contrib.auth.models.Permission" - reveal_type(User.groups) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyToManyDescriptor[django.db.models.base.Model]" - reveal_type(User.user_permissions) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyToManyDescriptor[django.db.models.base.Model]" + reveal_type(User.groups) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyToManyDescriptor[django.contrib.auth.models.Group, django.db.models.base.Model]" + reveal_type(User.user_permissions) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyToManyDescriptor[django.contrib.auth.models.Permission, django.db.models.base.Model]" from django.contrib.auth.models import AnonymousUser reveal_type(AnonymousUser().is_authenticated) # N: Revealed type is "Literal[False]" @@ -33,7 +33,7 @@ from django.contrib.auth.models import Group reveal_type(Group().name) # N: Revealed type is "builtins.str" reveal_type(Group().permissions.get()) # N: Revealed type is "django.contrib.auth.models.Permission" - reveal_type(Group.permissions) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyToManyDescriptor[django.contrib.auth.models.Group_permissions]" + reveal_type(Group.permissions) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyToManyDescriptor[django.contrib.auth.models.Permission, django.contrib.auth.models.Group_permissions]" - case: can_override_abstract_user_manager main: | From f1425aa6286ad9655f61219130bc6db8a56e4c9b Mon Sep 17 00:00:00 2001 From: Marti Raudsepp Date: Fri, 10 Nov 2023 17:54:47 +0200 Subject: [PATCH 2/7] mypy 1.7 fix --- tests/typecheck/fields/test_related.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/typecheck/fields/test_related.yml b/tests/typecheck/fields/test_related.yml index a0c4034b4..0a1cdf3b9 100644 --- a/tests/typecheck/fields/test_related.yml +++ b/tests/typecheck/fields/test_related.yml @@ -1134,7 +1134,7 @@ main:26: note: Revealed type is "myapp.models.MyModel_auto_through" main:28: note: Revealed type is "builtins.str" main:29: note: Revealed type is "builtins.str" - main:30: note: Revealed type is "def (*objs: Union[myapp.models.MyModel, builtins.int], *, bulk: builtins.bool =)" + main:30: note: Revealed type is "def (*objs: Union[myapp.models.MyModel, builtins.int], bulk: builtins.bool =)" installed_apps: - myapp files: From a6cc98daafd89107a09ccb34fecd2b2f75d0f56a Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Thu, 16 Nov 2023 09:50:19 +0100 Subject: [PATCH 3/7] fixup! Improve hints for `ManyToManyDescriptor` and start using it --- django-stubs/db/models/fields/related_descriptors.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/django-stubs/db/models/fields/related_descriptors.pyi b/django-stubs/db/models/fields/related_descriptors.pyi index bb3f60d15..580144e53 100644 --- a/django-stubs/db/models/fields/related_descriptors.pyi +++ b/django-stubs/db/models/fields/related_descriptors.pyi @@ -141,7 +141,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) -> Self: ... + def __call__(self, *, manager: str) -> ManyRelatedManager[_M]: ... def create_forward_many_to_many_manager( superclass: type[BaseManager[_M]], rel: ManyToManyRel, reverse: bool From b87f9dc0dad588294c0dfc5694509395be8f6819 Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Thu, 16 Nov 2023 09:58:24 +0100 Subject: [PATCH 4/7] fixup! fixup! Improve hints for `ManyToManyDescriptor` and start using it --- tests/typecheck/fields/test_related.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/typecheck/fields/test_related.yml b/tests/typecheck/fields/test_related.yml index a23fdcc18..9b49568a7 100644 --- a/tests/typecheck/fields/test_related.yml +++ b/tests/typecheck/fields/test_related.yml @@ -946,7 +946,7 @@ from myapp.models import SalesMan sales_man = SalesMan() reveal_type(sales_man.client) # N: Revealed type is "myapp.models.CustomUser_ManyRelatedManager" - reveal_type(sales_man.client(manager="staffs")) # N: Revealed type is "myapp.models.CustomUser_ManyRelatedManager" + reveal_type(sales_man.client(manager="staffs")) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyRelatedManager[myapp.models.CustomUser]" installed_apps: - myapp files: From 866645a33fdb34cbc945f049d0d6cbde89ac3a68 Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Mon, 27 Nov 2023 16:44:30 +0100 Subject: [PATCH 5/7] fixup! fixup! fixup! Improve hints for `ManyToManyDescriptor` and start using it --- mypy_django_plugin/lib/helpers.py | 7 +++ mypy_django_plugin/main.py | 1 + mypy_django_plugin/transformers/manytomany.py | 43 +++++++++++-------- mypy_django_plugin/transformers/models.py | 4 +- 4 files changed, 35 insertions(+), 20 deletions(-) diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py index 4d1da8c6c..97c2c8f2f 100644 --- a/mypy_django_plugin/lib/helpers.py +++ b/mypy_django_plugin/lib/helpers.py @@ -471,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 + ) diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py index 19ed7c743..ec6bb3aed 100644 --- a/mypy_django_plugin/main.py +++ b/mypy_django_plugin/main.py @@ -187,6 +187,7 @@ 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, diff --git a/mypy_django_plugin/transformers/manytomany.py b/mypy_django_plugin/transformers/manytomany.py index 1f3b4715c..e98f91153 100644 --- a/mypy_django_plugin/transformers/manytomany.py +++ b/mypy_django_plugin/transformers/manytomany.py @@ -153,33 +153,38 @@ def get_model_from_expression( 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. """ - 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 - ): + 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 = many_related_manager.args[0].copy_modified() + 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" ) diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index 4920e67be..d6f7f910d 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -484,8 +484,9 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: ), ) continue + elif isinstance(relation, ManyToManyRel): - # TODO: 'relation' should be based on `TypeInfo` instead of Django runtime.. + # 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 @@ -497,6 +498,7 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: self.many_to_many_descriptor, [Instance(to_model_info, []), Instance(through_model_info, [])] ), ) + else: related_manager_info = None try: From e5462f185c031fae725071119bb0921638d83141 Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Mon, 27 Nov 2023 16:47:41 +0100 Subject: [PATCH 6/7] fixup! fixup! fixup! fixup! Improve hints for `ManyToManyDescriptor` and start using it --- django-stubs/db/models/fields/related.pyi | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/django-stubs/db/models/fields/related.pyi b/django-stubs/db/models/fields/related.pyi index 67525ac7e..2c60b595e 100644 --- a/django-stubs/db/models/fields/related.pyi +++ b/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[_To, _M]: ... + def __get__(self, instance: None, owner: Any) -> ManyToManyDescriptor[_To, _Through]: ... # Model instance access @overload def __get__(self, instance: Model, owner: Any) -> ManyRelatedManager[_To]: ... From d45a3db5600fdf423abdbd4c5c37366ab4c73433 Mon Sep 17 00:00:00 2001 From: Marti Raudsepp Date: Mon, 27 Nov 2023 20:15:45 +0200 Subject: [PATCH 7/7] Fix Python 3.8 --- mypy_django_plugin/transformers/manytomany.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mypy_django_plugin/transformers/manytomany.py b/mypy_django_plugin/transformers/manytomany.py index e98f91153..494174a1a 100644 --- a/mypy_django_plugin/transformers/manytomany.py +++ b/mypy_django_plugin/transformers/manytomany.py @@ -1,4 +1,4 @@ -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 @@ -153,7 +153,7 @@ 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]]: if ( isinstance(ctx.default_return_type, Instance) and ctx.default_return_type.type.fullname == fullnames.MANY_RELATED_MANAGER