From 06422a51e70ff62aafe6418ff0b4032adee76e1a Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Fri, 22 Mar 2024 20:05:20 +0100 Subject: [PATCH] fixup! Add a generic type parameter 'through' to `ManyRelatedManager` --- django-stubs/db/models/fields/related_descriptors.pyi | 2 +- ext/django_stubs_ext/db/models/manager.py | 4 ++-- mypy_django_plugin/transformers/manytomany.py | 11 +++++++++++ tests/typecheck/fields/test_related.yml | 8 ++++---- tests/typecheck/models/test_contrib_models.yml | 4 ++-- 5 files changed, 20 insertions(+), 9 deletions(-) diff --git a/django-stubs/db/models/fields/related_descriptors.pyi b/django-stubs/db/models/fields/related_descriptors.pyi index 1b60c0d489..c3745e2b5b 100644 --- a/django-stubs/db/models/fields/related_descriptors.pyi +++ b/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): diff --git a/ext/django_stubs_ext/db/models/manager.py b/ext/django_stubs_ext/db/models/manager.py index 81f340166a..4e7684b08e 100644 --- a/ext/django_stubs_ext/db/models/manager.py +++ b/ext/django_stubs_ext/db/models/manager.py @@ -13,12 +13,12 @@ from typing import Protocol, TypeVar _T = TypeVar("_T") - _M = TypeVar("_M") + _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, _M]): + class ManyRelatedManager(Protocol[_T, _Through]): pass diff --git a/mypy_django_plugin/transformers/manytomany.py b/mypy_django_plugin/transformers/manytomany.py index b8ab86f038..2de97526d9 100644 --- a/mypy_django_plugin/transformers/manytomany.py +++ b/mypy_django_plugin/transformers/manytomany.py @@ -151,6 +151,17 @@ def get_model_from_expression( 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 instance of 1. when it's a model + 3. The second type parameter 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 diff --git a/tests/typecheck/fields/test_related.yml b/tests/typecheck/fields/test_related.yml index 3d048bdaa0..861c58ef32 100644 --- a/tests/typecheck/fields/test_related.yml +++ b/tests/typecheck/fields/test_related.yml @@ -603,7 +603,7 @@ 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.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 "myapp.models.Blog_ManyRelatedManager" + reveal_type(Author().blogs) # N: Revealed type is "myapp.models.Blog_ManyRelatedManager[myapp.models.Author_blogs]" 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" @@ -946,8 +946,8 @@ main: | 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 "django.db.models.fields.related_descriptors.ManyRelatedManager[myapp.models.CustomUser]" + reveal_type(sales_man.client) # N: Revealed type is "myapp.models.CustomUser_ManyRelatedManager[myapp.models.SalesMan_client]" + reveal_type(sales_man.client(manager="staffs")) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyRelatedManager[myapp.models.CustomUser, myapp.models.SalesMan_client]" installed_apps: - myapp files: @@ -1279,7 +1279,7 @@ instance = arg(field=1, unknown=2) reveal_type(instance.field) # N: Revealed type is "Any" reveal_type(instance.unknown) # N: Revealed type is "Any" - reveal_type(instance.bad) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyRelatedManager[Any]" + reveal_type(instance.bad) # N: Revealed type is "django.db.models.fields.related_descriptors.ManyRelatedManager[Any, Any]" return instance installed_apps: - myapp diff --git a/tests/typecheck/models/test_contrib_models.yml b/tests/typecheck/models/test_contrib_models.yml index 717044bc31..16d4910d80 100644 --- a/tests/typecheck/models/test_contrib_models.yml +++ b/tests/typecheck/models/test_contrib_models.yml @@ -163,10 +163,10 @@ - case: test_permissions_inherited_reverese_relations main: | from django.contrib.auth.models import Permission - reveal_type(Permission().user_set) # N: Revealed type is "django.contrib.auth.models.User_ManyRelatedManager" + reveal_type(Permission().user_set) # N: Revealed type is "django.contrib.auth.models.User_ManyRelatedManager[django.contrib.auth.models.User_user_permissions]" from django.contrib.auth.models import Group - reveal_type(Group().user_set) # N: Revealed type is "django.contrib.auth.models.User_ManyRelatedManager" + reveal_type(Group().user_set) # N: Revealed type is "django.contrib.auth.models.User_ManyRelatedManager[django.contrib.auth.models.User_groups]" custom_settings: | INSTALLED_APPS = ('django.contrib.contenttypes', 'django.contrib.auth') AUTH_USER_MODEL='auth.User'