From 60b96e60ebee459369e6d8041e02e4d6e3650828 Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Tue, 5 Sep 2023 08:41:36 +0200 Subject: [PATCH] Only populate a model argument for generic managers --- mypy_django_plugin/transformers/models.py | 17 ++- .../managers/querysets/test_from_queryset.yml | 103 ++++++++++++++++-- tests/typecheck/managers/test_managers.yml | 38 +++++-- .../typecheck/models/test_contrib_models.yml | 2 +- 4 files changed, 139 insertions(+), 21 deletions(-) diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index 9c41ba961..7a031718d 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -314,8 +314,11 @@ def reparametrize_dynamically_created_manager(self, manager_name: str, manager_i return 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(): + # Reparameterize dynamically created manager with model type + manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])]) + else: + manager_type = Instance(manager_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: @@ -340,7 +343,10 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None: incomplete_manager_defs.add(manager_name) continue - manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])]) + if manager_info.is_generic(): + manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])]) + else: + manager_type = Instance(manager_info, []) self.add_new_node_to_model_class(manager_name, manager_type, is_classvar=True) if incomplete_manager_defs: @@ -439,7 +445,10 @@ 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(): + default_manager = Instance(default_manager_info, [Instance(self.model_classdef.info, [])]) + else: + default_manager = Instance(default_manager_info, []) self.add_new_node_to_model_class("_default_manager", default_manager, is_classvar=True) diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index 59f0cd63c..dd28e33f2 100644 --- a/tests/typecheck/managers/querysets/test_from_queryset.yml +++ b/tests/typecheck/managers/querysets/test_from_queryset.yml @@ -438,14 +438,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.all) # N: Revealed type is "def () -> myapp.models.MyQuerySet[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[myapp.models.MyModel]" + 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: @@ -468,7 +468,7 @@ - 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" installed_apps: - myapp files: @@ -675,3 +675,92 @@ ... StrCallable = BaseManager.from_queryset(ModelQuerySet, class_name=str(1)) + +- case: only_auto_populates_model_arg_to_manager_with_type_var + main: | + from myapp.models import MyModel + reveal_type(MyModel.bm_populated_from_qs_populated) + reveal_type(MyModel.bm_generic_from_qs_populated) + + reveal_type(MyModel.bm_populated_from_qs_generic) + reveal_type(MyModel.bm_generic_from_qs_generic) + + reveal_type(MyModel.m_populated_from_qs_populated) + reveal_type(MyModel.m_generic_from_qs_populated) + + reveal_type(MyModel.m_populated_from_qs_generic) + reveal_type(MyModel.m_generic_from_qs_generic) + + reveal_type(MyModel.qs_populated_as_manager) + reveal_type(MyModel.qs_generic_as_manager) + out: | + main:2: note: Revealed type is "myapp.models.BaseManagerPopulatedFromQuerySetPopulated" + main:3: note: Revealed type is "myapp.models.BaseManagerGenericFromQuerySetPopulated[myapp.models.MyModel]" + main:5: note: Revealed type is "myapp.models.BaseManagerPopulatedFromQuerySetGeneric" + main:6: note: Revealed type is "myapp.models.BaseManagerGenericFromQuerySetGeneric[myapp.models.MyModel]" + main:8: note: Revealed type is "myapp.models.ManagerPopulatedFromQuerySetPopulated" + main:9: note: Revealed type is "myapp.models.ManagerGenericFromQuerySetPopulated[myapp.models.MyModel]" + main:11: note: Revealed type is "myapp.models.ManagerPopulatedFromQuerySetGeneric" + main:12: note: Revealed type is "myapp.models.ManagerGenericFromQuerySetGeneric[myapp.models.MyModel]" + main:14: note: Revealed type is "myapp.models.ManagerFromQuerySetPopulated[myapp.models.MyModel]" + # TODO: We want a line like below and not an error.. But mypy doesn't trigger + # the dynamic class hook for classes with generic arguments.. + # Revealed type is "myapp.models.ManagerFromQuerySetGeneric[myapp.models.MyModel]" + main:15: note: Revealed type is "myapp.models.UnknownManager[myapp.models.MyModel]" + myapp/models:50: error: Could not resolve manager type for "myapp.models.MyModel.qs_generic_as_manager" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from typing import TypeVar + from django.db.models import Manager, Model, QuerySet + from django.db.models.manager import BaseManager + + class QuerySetPopulated(QuerySet["MyModel"]): + ... + + T = TypeVar("T", bound=Model) + class QuerySetGeneric(QuerySet[T]): + ... + + class BaseManagerPopulated(BaseManager["MyModel"]): + ... + + class BaseManagerGeneric(BaseManager[T]): + ... + + class ManagerPopulated(Manager["MyModel"]): + ... + + class ManagerGeneric(Manager[T]): + ... + + BaseManagerPopulatedFromQuerySetPopulated = BaseManagerPopulated.from_queryset(QuerySetPopulated) + BaseManagerGenericFromQuerySetPopulated = BaseManagerGeneric.from_queryset(QuerySetPopulated) + + BaseManagerPopulatedFromQuerySetGeneric = BaseManagerPopulated.from_queryset(QuerySetGeneric) + BaseManagerGenericFromQuerySetGeneric = BaseManagerGeneric.from_queryset(QuerySetGeneric) + + ManagerPopulatedFromQuerySetPopulated = ManagerPopulated.from_queryset(QuerySetPopulated) + ManagerGenericFromQuerySetPopulated = ManagerGeneric.from_queryset(QuerySetPopulated) + + ManagerPopulatedFromQuerySetGeneric = ManagerPopulated.from_queryset(QuerySetGeneric) + ManagerGenericFromQuerySetGeneric = ManagerGeneric.from_queryset(QuerySetGeneric) + + class MyModel(Model): + bm_populated_from_qs_populated = BaseManagerPopulatedFromQuerySetPopulated() + bm_generic_from_qs_populated = BaseManagerGenericFromQuerySetPopulated() + + bm_populated_from_qs_generic = BaseManagerPopulatedFromQuerySetGeneric() + bm_generic_from_qs_generic = BaseManagerGenericFromQuerySetGeneric() + + m_populated_from_qs_populated = ManagerPopulatedFromQuerySetPopulated() + m_generic_from_qs_populated = ManagerGenericFromQuerySetPopulated() + + m_populated_from_qs_generic = ManagerPopulatedFromQuerySetGeneric() + m_generic_from_qs_generic = ManagerGenericFromQuerySetGeneric() + + qs_populated_as_manager = QuerySetPopulated.as_manager() + qs_generic_as_manager = QuerySetGeneric["MyModel"].as_manager() diff --git a/tests/typecheck/managers/test_managers.yml b/tests/typecheck/managers/test_managers.yml index 72c79b2e5..4a996be37 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 @@ -196,20 +196,33 @@ main: | from myapp.models import AbstractBase1, AbstractBase2, Child reveal_type(Child.manager1) + reveal_type(Child.manager1.get()) reveal_type(Child.restricted) + reveal_type(Child.restricted.get()) + reveal_type(Child.manager3) + reveal_type(Child.manager3.get()) reveal_type(AbstractBase1.manager1) + reveal_type(AbstractBase1.manager1.get()) reveal_type(AbstractBase2.restricted) + reveal_type(AbstractBase2.restricted.get()) 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.AbstractBase1" + main:4: note: Revealed type is "myapp.models.CustomManager2" + main:5: note: Revealed type is "myapp.models.AbstractBase2" + main:6: note: Revealed type is "myapp.models.GenericManager[myapp.models.Child]" + main:7: note: Revealed type is "myapp.models.Child" + main:8: note: Revealed type is "myapp.models.CustomManager1" + main:9: note: Revealed type is "myapp.models.AbstractBase1" + main:10: note: Revealed type is "myapp.models.CustomManager2" + main:11: note: Revealed type is "myapp.models.AbstractBase2" installed_apps: - myapp files: - path: myapp/__init__.py - path: myapp/models.py content: | + from typing import TypeVar from django.db import models class CustomManager1(models.Manager['AbstractBase1']): pass @@ -225,8 +238,15 @@ abstract = True value = models.CharField(max_length=50) restricted = CustomManager2() + T = TypeVar("T", bound="AbstractBase3") + class GenericManager(models.Manager[T]): + ... + class AbstractBase3(models.Model): + class Meta: + abstract = True + manager3 = GenericManager() - class Child(AbstractBase1, AbstractBase2): + class Child(AbstractBase1, AbstractBase2, AbstractBase3): pass - case: model_has_a_manager_of_the_same_type @@ -639,7 +659,7 @@ 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" files: - path: myapp/__init__.py - path: myapp/models.py diff --git a/tests/typecheck/models/test_contrib_models.yml b/tests/typecheck/models/test_contrib_models.yml index 9c8ed4f10..9326b90e3 100644 --- a/tests/typecheck/models/test_contrib_models.yml +++ b/tests/typecheck/models/test_contrib_models.yml @@ -32,7 +32,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]"