Skip to content

Commit

Permalink
Don't populate a type arg for non generic managers (#2314)
Browse files Browse the repository at this point in the history
  • Loading branch information
flaeppe committed Aug 3, 2024
1 parent 2f4c99a commit ef501f2
Show file tree
Hide file tree
Showing 6 changed files with 118 additions and 22 deletions.
4 changes: 4 additions & 0 deletions mypy_django_plugin/lib/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 [])
13 changes: 5 additions & 8 deletions mypy_django_plugin/transformers/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -318,7 +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, [])])
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:
Expand All @@ -344,7 +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])
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:
Expand All @@ -361,11 +361,8 @@ 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 = 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()`
manager_expr = self.get_manager_expression(manager_name)
Expand Down Expand Up @@ -445,7 +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, [])])
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)


Expand Down
1 change: 1 addition & 0 deletions mypy_django_plugin/transformers/querysets.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
13 changes: 8 additions & 5 deletions tests/typecheck/managers/querysets/test_from_queryset.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand Down
107 changes: 99 additions & 8 deletions tests/typecheck/managers/test_managers.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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()
2 changes: 1 addition & 1 deletion tests/typecheck/models/test_contrib_models.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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]"
Expand Down

0 comments on commit ef501f2

Please sign in to comment.