Only populate a model argument for generic managers

Author: flaeppe

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)
 
 

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()

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

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]"