From edccb7719a4817d889a2d6d1b492ca9fdfc98b53 Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Sat, 3 Aug 2024 11:31:09 +0200 Subject: [PATCH 1/2] Don't populate a type arg for non generic managers --- mypy_django_plugin/transformers/models.py | 20 ++-- mypy_django_plugin/transformers/querysets.py | 1 + .../managers/querysets/test_from_queryset.yml | 13 ++- tests/typecheck/managers/test_managers.yml | 107 ++++++++++++++++-- .../typecheck/models/test_contrib_models.yml | 2 +- 5 files changed, 122 insertions(+), 21 deletions(-) diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index b77dc7dec..2612efee0 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -318,7 +318,9 @@ def reparametrize_dynamically_created_manager(self, manager_name: str, manager_i assert manager_info is not None # Reparameterize dynamically created manager with model type - manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])]) + manager_type = Instance( + manager_info, [Instance(self.model_classdef.info, [])] if manager_info.is_generic() else [] + ) self.add_new_node_to_model_class(manager_name, manager_type, is_classvar=True) def run_with_model_cls(self, model_cls: Type[Model]) -> None: @@ -344,7 +346,9 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: continue assert self.model_classdef.info.self_type is not None - manager_type = Instance(manager_info, [self.model_classdef.info.self_type]) + manager_type = Instance( + manager_info, [self.model_classdef.info.self_type] if manager_info.is_generic() else [] + ) self.add_new_node_to_model_class(manager_name, manager_type, is_classvar=True) if incomplete_manager_defs: @@ -361,11 +365,11 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: fallback_manager_info = self.get_or_create_manager_with_any_fallback() if fallback_manager_info is not None: assert self.model_classdef.info.self_type is not None - self.add_new_node_to_model_class( - manager_name, - Instance(fallback_manager_info, [self.model_classdef.info.self_type]), - is_classvar=True, + manager_type = Instance( + fallback_manager_info, + [self.model_classdef.info.self_type] if fallback_manager_info.is_generic() else [], ) + self.add_new_node_to_model_class(manager_name, manager_type, is_classvar=True) # Find expression for e.g. `objects = SomeManager()` manager_expr = self.get_manager_expression(manager_name) @@ -445,7 +449,9 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: return None default_manager_info = generated_manager_info - default_manager = Instance(default_manager_info, [Instance(self.model_classdef.info, [])]) + default_manager = Instance( + default_manager_info, [Instance(self.model_classdef.info, [])] if default_manager_info.is_generic() else [] + ) self.add_new_node_to_model_class("_default_manager", default_manager, is_classvar=True) diff --git a/mypy_django_plugin/transformers/querysets.py b/mypy_django_plugin/transformers/querysets.py index 1fb4e2a6a..0ca9a76a5 100644 --- a/mypy_django_plugin/transformers/querysets.py +++ b/mypy_django_plugin/transformers/querysets.py @@ -48,6 +48,7 @@ def determine_proper_manager_type(ctx: FunctionContext) -> MypyType: outer_model_info is None or not outer_model_info.has_base(fullnames.MODEL_CLASS_FULLNAME) or outer_model_info.self_type is None + or not default_return_type.type.is_generic() ): return default_return_type diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index ee73505eb..60f6c765f 100644 --- a/tests/typecheck/managers/querysets/test_from_queryset.yml +++ b/tests/typecheck/managers/querysets/test_from_queryset.yml @@ -583,14 +583,14 @@ - case: test_queryset_in_model_class_body main: | from myapp.models import MyModel - reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]" - reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet" + reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet" reveal_type(MyModel.objects.all) # N: Revealed type is "def () -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.custom) # N: Revealed type is "def () -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.all().filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.custom().filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet" - reveal_type(MyModel.objects2) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]" - reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects2) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet" + reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet" installed_apps: - myapp files: @@ -613,7 +613,10 @@ - case: test_queryset_in_model_class_body_subclass main: | from myapp.models import MyModel - reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.BaseManagerFromBaseQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.BaseManagerFromBaseQuerySet" + reveal_type(MyModel.objects.get()) # N: Revealed type is "myapp.models.BaseModel" + reveal_type(MyModel.objects.filter()) # N: Revealed type is "myapp.models.BaseQuerySet" + reveal_type(MyModel.objects.filter().get()) # N: Revealed type is "myapp.models.BaseModel" installed_apps: - myapp files: diff --git a/tests/typecheck/managers/test_managers.yml b/tests/typecheck/managers/test_managers.yml index fd360c727..1a9b0019f 100644 --- a/tests/typecheck/managers/test_managers.yml +++ b/tests/typecheck/managers/test_managers.yml @@ -126,7 +126,7 @@ - case: test_leave_as_is_if_objects_is_set_and_fill_typevars_with_outer_class main: | from myapp.models import MyUser - reveal_type(MyUser.objects) # N: Revealed type is "myapp.models.UserManager[myapp.models.MyUser]" + reveal_type(MyUser.objects) # N: Revealed type is "myapp.models.UserManager" reveal_type(MyUser.objects.get()) # N: Revealed type is "myapp.models.MyUser" reveal_type(MyUser.objects.get_or_404()) # N: Revealed type is "myapp.models.MyUser" installed_apps: @@ -169,9 +169,9 @@ main: | from myapp.models import AbstractPerson, Book reveal_type(AbstractPerson.abstract_persons) # N: Revealed type is "django.db.models.manager.Manager[myapp.models.AbstractPerson]" - reveal_type(Book.published_objects) # N: Revealed type is "myapp.models.PublishedBookManager[myapp.models.Book]" + reveal_type(Book.published_objects) # N: Revealed type is "myapp.models.PublishedBookManager" Book.published_objects.create(title='hello') - reveal_type(Book.annotated_objects) # N: Revealed type is "myapp.models.AnnotatedBookManager[myapp.models.Book]" + reveal_type(Book.annotated_objects) # N: Revealed type is "myapp.models.AnnotatedBookManager" Book.annotated_objects.create(title='hello') installed_apps: - myapp @@ -200,10 +200,10 @@ reveal_type(AbstractBase1.manager1) reveal_type(AbstractBase2.restricted) out: | - main:2: note: Revealed type is "myapp.models.CustomManager1[myapp.models.Child]" - main:3: note: Revealed type is "myapp.models.CustomManager2[myapp.models.Child]" - main:4: note: Revealed type is "myapp.models.CustomManager1[myapp.models.AbstractBase1]" - main:5: note: Revealed type is "myapp.models.CustomManager2[myapp.models.AbstractBase2]" + main:2: note: Revealed type is "myapp.models.CustomManager1" + main:3: note: Revealed type is "myapp.models.CustomManager2" + main:4: note: Revealed type is "myapp.models.CustomManager1" + main:5: note: Revealed type is "myapp.models.CustomManager2" installed_apps: - myapp files: @@ -229,6 +229,44 @@ class Child(AbstractBase1, AbstractBase2): pass +- case: managers_inherited_from_abstract_classes_multiple_inheritance_with_generic + main: | + from myapp.models import AbstractBase1, AbstractBase2, Child + reveal_type(Child.manager1) # N: Revealed type is "myapp.models.CustomManager1[myapp.models.Child]" + reveal_type(Child.manager1.get()) # N: Revealed type is "myapp.models.Child" + reveal_type(Child.restricted) # N: Revealed type is "myapp.models.CustomManager2[myapp.models.Child]" + reveal_type(Child.restricted.get()) # N: Revealed type is "myapp.models.Child" + reveal_type(AbstractBase1.manager1) # N: Revealed type is "myapp.models.CustomManager1[myapp.models.AbstractBase1]" + reveal_type(AbstractBase1.manager1.get()) # N: Revealed type is "myapp.models.AbstractBase1" + reveal_type(AbstractBase2.restricted) # N: Revealed type is "myapp.models.CustomManager2[myapp.models.AbstractBase2]" + reveal_type(AbstractBase2.restricted.get()) # N: Revealed type is "myapp.models.AbstractBase2" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from typing import Generic, TypeVar + from django.db import models + + T1 = TypeVar("T1", bound="AbstractBase1") + class CustomManager1(models.Manager[T1], Generic[T1]): ... + class AbstractBase1(models.Model): + class Meta: + abstract = True + name = models.CharField(max_length=50) + manager1 = CustomManager1() + + T2 = TypeVar("T2", bound="AbstractBase2") + class CustomManager2(models.Manager[T2], Generic[T2]): ... + class AbstractBase2(models.Model): + class Meta: + abstract = True + value = models.CharField(max_length=50) + restricted = CustomManager2() + + class Child(AbstractBase1, AbstractBase2): ... + - case: model_has_a_manager_of_the_same_type main: | from myapp.models import UnrelatedModel, MyModel @@ -638,10 +676,12 @@ main: | from myapp.models import MyModel reveal_type(MyModel.objects) + reveal_type(MyModel.objects.get()) installed_apps: - myapp out: | - main:2: note: Revealed type is "myapp.models.MyModel.MyManager[myapp.models.MyModel]" + main:2: note: Revealed type is "myapp.models.MyModel.MyManager" + main:3: note: Revealed type is "Any" files: - path: myapp/__init__.py - path: myapp/models.py @@ -670,3 +710,54 @@ from django.db import models class MyModel(models.Model): ... + +- case: test_does_not_populate_an_unexpected_type_argument + main: | + from myapp.models import MyModel + reveal_type(MyModel.populated_manager) # N: Revealed type is "myapp.models.PopulatedManager" + reveal_type(MyModel.populated_manager.get()) # N: Revealed type is "myapp.models.MyModel" + reveal_type(MyModel.populated_manager.filter()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel, myapp.models.MyModel]" + + reveal_type(MyModel.populated_manager_from_generic_queryset) # N: Revealed type is "myapp.models.PopulatedManagerFromQuerySet" + reveal_type(MyModel.populated_manager_from_generic_queryset.get()) # N: Revealed type is "myapp.models.MyModel" + reveal_type(MyModel.populated_manager_from_generic_queryset.filter()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel, myapp.models.MyModel]" + + reveal_type(MyModel.populated_manager_from_populated_queryset) # N: Revealed type is "myapp.models.PopulatedManagerFromPopulatedQuerySet" + reveal_type(MyModel.populated_manager_from_populated_queryset.get()) # N: Revealed type is "myapp.models.MyModel" + reveal_type(MyModel.populated_manager_from_populated_queryset.filter()) # N: Revealed type is "myapp.models.PopulatedQuerySet" + + reveal_type(MyModel.generic_manager) # N: Revealed type is "django.db.models.manager.Manager[myapp.models.MyModel]" + reveal_type(MyModel.generic_manager.get()) # N: Revealed type is "myapp.models.MyModel" + reveal_type(MyModel.generic_manager.filter()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel, myapp.models.MyModel]" + + reveal_type(MyModel.generic_manager_from_generic_queryset) # N: Revealed type is "myapp.models.ManagerFromQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.generic_manager_from_generic_queryset.get()) # N: Revealed type is "myapp.models.MyModel" + reveal_type(MyModel.generic_manager_from_generic_queryset.filter()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyModel, myapp.models.MyModel]" + + reveal_type(MyModel.generic_manager_from_populated_queryset) # N: Revealed type is "myapp.models.ManagerFromPopulatedQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.generic_manager_from_populated_queryset.get()) # N: Revealed type is "myapp.models.MyModel" + reveal_type(MyModel.generic_manager_from_populated_queryset.filter()) # N: Revealed type is "myapp.models.PopulatedQuerySet" + + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + + class PopulatedManager(models.Manager["MyModel"]): ... + class PopulatedQuerySet(models.QuerySet["MyModel"]): ... + + PopulatedManagerFromGenericQuerySet = PopulatedManager.from_queryset(models.QuerySet) + PopulatedManagerFromPopulatedQuerySet = PopulatedManager.from_queryset(PopulatedQuerySet) + GenericManagerFromGenericQuerySet = models.Manager.from_queryset(models.QuerySet) + GenericManagerFromPopulatedQuerySet = models.Manager.from_queryset(PopulatedQuerySet) + + class MyModel(models.Model): + populated_manager = PopulatedManager() + populated_manager_from_generic_queryset = PopulatedManagerFromGenericQuerySet() + populated_manager_from_populated_queryset = PopulatedManagerFromPopulatedQuerySet() + generic_manager = models.Manager() + generic_manager_from_generic_queryset = GenericManagerFromGenericQuerySet() + generic_manager_from_populated_queryset = GenericManagerFromPopulatedQuerySet() diff --git a/tests/typecheck/models/test_contrib_models.yml b/tests/typecheck/models/test_contrib_models.yml index ca7b14be0..07165a8b2 100644 --- a/tests/typecheck/models/test_contrib_models.yml +++ b/tests/typecheck/models/test_contrib_models.yml @@ -38,7 +38,7 @@ - case: can_override_abstract_user_manager main: | from myapp.models import MyBaseUser, MyUser - reveal_type(MyBaseUser.objects) # N: Revealed type is "myapp.models.MyBaseUserManager[myapp.models.MyBaseUser]" + reveal_type(MyBaseUser.objects) # N: Revealed type is "myapp.models.MyBaseUserManager" reveal_type(MyBaseUser.objects.all()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyBaseUser, myapp.models.MyBaseUser]" reveal_type(MyUser.objects) # N: Revealed type is "myapp.models.MyUserManager" reveal_type(MyUser.objects.all()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.MyUser, myapp.models.MyUser]" From 776c863bd34b3dd26db9ad0a4f7bcc491262c601 Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Sat, 3 Aug 2024 12:34:04 +0200 Subject: [PATCH 2/2] fixup! Don't populate a type arg for non generic managers --- mypy_django_plugin/lib/helpers.py | 4 ++++ mypy_django_plugin/transformers/models.py | 17 ++++------------- 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py index 2d7d5bd80..fabcb68e1 100644 --- a/mypy_django_plugin/lib/helpers.py +++ b/mypy_django_plugin/lib/helpers.py @@ -533,3 +533,7 @@ def get_model_from_expression( if model_info is not None: return Instance(model_info, []) return None + + +def fill_manager(manager: TypeInfo, typ: MypyType) -> Instance: + return Instance(manager, [typ] if manager.is_generic() else []) diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index 2612efee0..92c3c8fc8 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -318,9 +318,7 @@ def reparametrize_dynamically_created_manager(self, manager_name: str, manager_i assert manager_info is not None # Reparameterize dynamically created manager with model type - manager_type = Instance( - manager_info, [Instance(self.model_classdef.info, [])] if manager_info.is_generic() else [] - ) + manager_type = helpers.fill_manager(manager_info, Instance(self.model_classdef.info, [])) self.add_new_node_to_model_class(manager_name, manager_type, is_classvar=True) def run_with_model_cls(self, model_cls: Type[Model]) -> None: @@ -346,9 +344,7 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: continue assert self.model_classdef.info.self_type is not None - manager_type = Instance( - manager_info, [self.model_classdef.info.self_type] if manager_info.is_generic() else [] - ) + manager_type = helpers.fill_manager(manager_info, self.model_classdef.info.self_type) self.add_new_node_to_model_class(manager_name, manager_type, is_classvar=True) if incomplete_manager_defs: @@ -365,10 +361,7 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: fallback_manager_info = self.get_or_create_manager_with_any_fallback() if fallback_manager_info is not None: assert self.model_classdef.info.self_type is not None - manager_type = Instance( - fallback_manager_info, - [self.model_classdef.info.self_type] if fallback_manager_info.is_generic() else [], - ) + manager_type = helpers.fill_manager(fallback_manager_info, self.model_classdef.info.self_type) self.add_new_node_to_model_class(manager_name, manager_type, is_classvar=True) # Find expression for e.g. `objects = SomeManager()` @@ -449,9 +442,7 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: return None default_manager_info = generated_manager_info - default_manager = Instance( - default_manager_info, [Instance(self.model_classdef.info, [])] if default_manager_info.is_generic() else [] - ) + default_manager = helpers.fill_manager(default_manager_info, Instance(self.model_classdef.info, [])) self.add_new_node_to_model_class("_default_manager", default_manager, is_classvar=True)