From 2f1ad70f94178c2a633bc6c0bef2df53eb7444f8 Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Fri, 22 Mar 2024 15:44:51 +0100 Subject: [PATCH] Add a generic type parameter 'through' to `ManyRelatedManager` `ManyRelatedManager` has a `.through` attribute which is typed with a generic type parameter. The plugin is also changed to populate and consider the new generic type parameter. --- django-stubs/db/models/fields/related.pyi | 2 +- .../db/models/fields/related_descriptors.pyi | 27 ++++---- ext/django_stubs_ext/db/models/manager.py | 3 +- mypy_django_plugin/lib/helpers.py | 15 +++++ mypy_django_plugin/transformers/manytomany.py | 61 ++++++++++++++----- tests/typecheck/fields/test_related.yml | 48 ++++++++++----- .../typecheck/models/test_contrib_models.yml | 4 +- 7 files changed, 115 insertions(+), 45 deletions(-) diff --git a/django-stubs/db/models/fields/related.pyi b/django-stubs/db/models/fields/related.pyi index ead3e49af..9c28014f0 100644 --- a/django-stubs/db/models/fields/related.pyi +++ b/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: ... diff --git a/django-stubs/db/models/fields/related_descriptors.pyi b/django-stubs/db/models/fields/related_descriptors.pyi index e5bda5cd1..c3745e2b5 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): @@ -125,23 +125,24 @@ 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 = ..., @@ -149,7 +150,7 @@ class ManyRelatedManager(Manager[_M], Generic[_M]): ) -> None: ... async def aset( self, - objs: QuerySet[_M] | Iterable[_M | int], + objs: QuerySet[_To] | Iterable[_To | int], *, bulk: bool = ..., clear: bool = ..., @@ -157,8 +158,8 @@ class ManyRelatedManager(Manager[_M], Generic[_M]): ) -> 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]]: ... diff --git a/ext/django_stubs_ext/db/models/manager.py b/ext/django_stubs_ext/db/models/manager.py index 97a4651b0..4e7684b08 100644 --- a/ext/django_stubs_ext/db/models/manager.py +++ b/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 diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py index abb8ac695..116097ac7 100644 --- a/mypy_django_plugin/lib/helpers.py +++ b/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 diff --git a/mypy_django_plugin/transformers/manytomany.py b/mypy_django_plugin/transformers/manytomany.py index d8467bc5f..94ce58d8f 100644 --- a/mypy_django_plugin/transformers/manytomany.py +++ b/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]) diff --git a/tests/typecheck/fields/test_related.yml b/tests/typecheck/fields/test_related.yml index dbd34c8b0..1f60c1653 100644 --- a/tests/typecheck/fields/test_related.yml +++ b/tests/typecheck/fields/test_related.yml @@ -278,11 +278,11 @@ - 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 "myapp.models.App_ManyRelatedManager" + reveal_type(Member().apps) # N: Revealed type is "myapp.models.App_ManyRelatedManager[myapp.models.Member_apps]" reveal_type(Member().apps.get()) # N: Revealed type is "myapp.models.App" 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) # N: Revealed type is "myapp.models.Member_ManyRelatedManager[myapp.models.Member_apps]" reveal_type(App().members.get()) # N: Revealed type is "myapp.models.Member" reveal_type(App().members.member_method()) # N: Revealed type is "builtins.int" App().members.add(Member()) @@ -311,7 +311,7 @@ - case: many_to_many_works_with_string_if_imported main: | from myapp.models import Member - reveal_type(Member().apps) # N: Revealed type is "myapp2.models.App_ManyRelatedManager" + reveal_type(Member().apps) # N: Revealed type is "myapp2.models.App_ManyRelatedManager[myapp.models.Member_apps]" installed_apps: - myapp - myapp2 @@ -346,7 +346,7 @@ - case: many_to_many_with_self main: | from myapp.models import User - reveal_type(User().friends) # N: Revealed type is "myapp.models.User_ManyRelatedManager" + reveal_type(User().friends) # N: Revealed type is "myapp.models.User_ManyRelatedManager[myapp.models.User_friends]" 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]" @@ -354,6 +354,7 @@ reveal_type(User.friends.through().from_user_id) # N: Revealed type is "builtins.int" reveal_type(User.friends.through().to_user) # N: Revealed type is "myapp.models.User" reveal_type(User.friends.through().to_user_id) # N: Revealed type is "builtins.int" + reveal_type(User.friends.through) # N: Revealed type is "Type[myapp.models.User_friends]" installed_apps: - myapp files: @@ -602,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" @@ -945,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: @@ -1114,6 +1115,16 @@ reveal_type(Other.autos.through.objects) reveal_type(Other.autos.through._default_manager) + + reveal_type(MyModel().auto_through) + reveal_type(MyModel().auto_through.through) + reveal_type(MyModel().auto_through.through.mymodel_id) + reveal_type(MyModel().other_again) + reveal_type(MyModel().other_again.through) + reveal_type(MyModel().other_again.through.mymodel_id) + reveal_type(MyModel().custom_through) + reveal_type(MyModel().custom_through.through) + reveal_type(MyModel().custom_through.through.my_model_id) out: | main:2: note: Revealed type is "myapp.models.MyModel_auto_through" main:3: note: Revealed type is "myapp.models.Other" @@ -1140,6 +1151,15 @@ main:30: note: Revealed type is "def (*objs: Union[myapp.models.MyModel, builtins.int], bulk: builtins.bool =, through_defaults: Union[builtins.dict[builtins.str, Any], None] =)" main:32: note: Revealed type is "django.db.models.manager.Manager[myapp.models.MyModel_auto_through]" main:33: note: Revealed type is "django.db.models.manager.Manager[myapp.models.MyModel_auto_through]" + main:35: note: Revealed type is "myapp.models.Other_ManyRelatedManager[myapp.models.MyModel_auto_through]" + main:36: note: Revealed type is "Type[myapp.models.MyModel_auto_through]" + main:37: note: Revealed type is "django.db.models.fields._FieldDescriptor[django.db.models.fields.AutoField[Union[django.db.models.expressions.Combinable, builtins.int, builtins.str, None], builtins.int]]" + main:38: note: Revealed type is "myapp.models.Other_ManyRelatedManager[myapp.models.MyModel_other_again]" + main:39: note: Revealed type is "Type[myapp.models.MyModel_other_again]" + main:40: note: Revealed type is "django.db.models.fields._FieldDescriptor[django.db.models.fields.AutoField[Union[django.db.models.expressions.Combinable, builtins.int, builtins.str, None], builtins.int]]" + main:41: note: Revealed type is "myapp.models.Other_ManyRelatedManager[myapp.models.CustomThrough]" + main:42: note: Revealed type is "Type[myapp.models.CustomThrough]" + main:43: note: Revealed type is "django.db.models.fields._FieldDescriptor[django.db.models.fields.AutoField[Union[django.db.models.expressions.Combinable, builtins.int, builtins.str], builtins.int]]" installed_apps: - myapp files: @@ -1231,7 +1251,7 @@ reveal_type(Child().parents) out: | 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]" + main:3: note: Revealed type is "django.db.models.fields.related_descriptors.ManyRelatedManager[Any, 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] installed_apps: @@ -1265,7 +1285,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 @@ -1284,11 +1304,11 @@ main: | from myapp.models import Other other = Other() - reveal_type(other.abstract_of_concrete_parent) # N: Revealed type is "myapp.models.ConcreteParent_ManyRelatedManager" - reveal_type(other.concrete_parent) # N: Revealed type is "myapp.models.ConcreteParent_ManyRelatedManager" - reveal_type(other.abstract_of_abstract_parent) # N: Revealed type is "myapp.models.MyModel_ManyRelatedManager" - reveal_type(other.abstract_parent) # N: Revealed type is "myapp.models.MyModel_ManyRelatedManager" - reveal_type(other.mymodel) # N: Revealed type is "myapp.models.MyModel_ManyRelatedManager" + reveal_type(other.abstract_of_concrete_parent) # N: Revealed type is "myapp.models.ConcreteParent_ManyRelatedManager[myapp.models.ConcreteParent_m2m_1]" + reveal_type(other.concrete_parent) # N: Revealed type is "myapp.models.ConcreteParent_ManyRelatedManager[myapp.models.ConcreteParent_m2m_2]" + reveal_type(other.abstract_of_abstract_parent) # N: Revealed type is "myapp.models.MyModel_ManyRelatedManager[myapp.models.MyModel_m2m_3]" + reveal_type(other.abstract_parent) # N: Revealed type is "myapp.models.MyModel_ManyRelatedManager[myapp.models.MyModel_m2m_4]" + reveal_type(other.mymodel) # N: Revealed type is "myapp.models.MyModel_ManyRelatedManager[myapp.models.MyModel_m2m_5]" installed_apps: - myapp files: diff --git a/tests/typecheck/models/test_contrib_models.yml b/tests/typecheck/models/test_contrib_models.yml index 717044bc3..16d4910d8 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'