fixup! Emit error when manager model argument is incorrect

mypy_django_plugin/transformers/managers.py

@@ -241,20 +241,15 @@ def merge_queryset_into_manager(
         )
 
 
-def get_dynamic_manager_model_argument_type(manager_info: TypeInfo, queryset_info: TypeInfo) -> Optional[MypyType]:
+def get_manager_model_argument_type(manager_info: TypeInfo) -> Optional[MypyType]:
     """
     TODO: Fill me out
     """
-    # TODO: Add test with manager inheritance
-    # TODO: Add test with queryset inheritance
     manager_model_arg = next(
         (
             base.args[0]
             for base in helpers.iter_bases(manager_info)
-            if (
-                base.type.fullname in {fullnames.MANAGER_CLASS_FULLNAME, fullnames.BASE_MANAGER_CLASS_FULLNAME}
-                and len(base.args) > 0
-            )
+            if (base.type.fullname in fullnames.MANAGER_CLASSES and len(base.args) > 0)
         ),
         (
             AnyType(TypeOfAny.from_omitted_generics)
@@ -264,10 +259,22 @@ def get_dynamic_manager_model_argument_type(manager_info: TypeInfo, queryset_inf
     )
     if isinstance(manager_model_arg, TypeVarType):
         # Since `django.db.models.Manager` is a subclass of `BaseManager` we'll collect
-        # its default `TypeVar` when argument has been omitted. Then we'll refine the
+        # its default `TypeVar` when argument has been omitted. So now we'll refine the
         # manager argument to be the type var's upper bound.
         manager_model_arg = manager_model_arg.upper_bound
 
+    return manager_model_arg
+
+
+def intersect_manager_and_queryset_model_argument_type(
+    manager_info: TypeInfo, queryset_info: TypeInfo
+) -> Optional[MypyType]:
+    """
+    TODO: Fill me out
+    """
+    # TODO: Add test with manager inheritance
+    # TODO: Add test with queryset inheritance
+    manager_model_arg = get_manager_model_argument_type(manager_info)
     queryset_model_arg = next(
         (
             base.args[0]
@@ -422,10 +429,10 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
         class_name=manager_class_name,
     )
 
-    model_argument_type = get_dynamic_manager_model_argument_type(manager_base, queryset_info)
+    model_argument_type = intersect_manager_and_queryset_model_argument_type(manager_base, queryset_info)
     if model_argument_type is None:
         semanal_api.fail(
-            "Manager model argument not matching queryset model argument",
+            "Manager model argument does not match queryset model argument",
             ctx.call,
             code=errorcodes.MODEL_ARG_MISMATCH,
         )

mypy_django_plugin/transformers/models.py

@@ -9,17 +9,21 @@
 from mypy.plugin import AnalyzeTypeContext, AttributeContext, CheckerPluginInterface, ClassDefContext
 from mypy.plugins import common
 from mypy.semanal import SemanticAnalyzer
-from mypy.types import AnyType, Instance
+from mypy.types import AnyType, Instance, PlaceholderType
 from mypy.types import Type as MypyType
 from mypy.types import TypedDictType, TypeOfAny
 
 from mypy_django_plugin.django.context import DjangoContext
-from mypy_django_plugin.errorcodes import MANAGER_MISSING
+from mypy_django_plugin.errorcodes import MANAGER_MISSING, MODEL_ARG_MISMATCH
 from mypy_django_plugin.lib import fullnames, helpers
 from mypy_django_plugin.lib.fullnames import ANNOTATIONS_FULLNAME, ANY_ATTR_ALLOWED_CLASS_FULLNAME, MODEL_CLASS_FULLNAME
 from mypy_django_plugin.lib.helpers import add_new_class_for_module
 from mypy_django_plugin.transformers import fields
 from mypy_django_plugin.transformers.fields import get_field_descriptor_types
+from mypy_django_plugin.transformers.managers import (
+    get_manager_model_argument_type,
+    intersect_manager_and_queryset_model_argument_type,
+)
 
 
 class ModelClassInitializer:
@@ -192,7 +196,11 @@ def has_any_parametrized_manager_as_base(self, info: TypeInfo) -> bool:
         return False
 
     def is_any_parametrized_manager(self, typ: Instance) -> bool:
-        return typ.type.fullname in fullnames.MANAGER_CLASSES and isinstance(typ.args[0], AnyType)
+        return (
+            typ.type.fullname in fullnames.MANAGER_CLASSES
+            and isinstance(typ.args[0], AnyType)
+            and typ.args[0].type_of_any == TypeOfAny.from_omitted_generics
+        )
 
     def create_new_model_parametrized_manager(self, name: str, base_manager_info: TypeInfo) -> Instance:
         bases = []
@@ -231,6 +239,50 @@ def create_new_model_parametrized_manager(self, name: str, base_manager_info: Ty
 
         return custom_manager_type
 
+    def manager_is_dynamically_generated(self, manager_info: TypeInfo) -> bool:
+        return manager_info.metadata.get("django", {}).get("from_queryset_manager") is not None
+
+    def get_manager_model_argument_type(self, manager_info: TypeInfo) -> Optional[MypyType]:
+        if self.manager_is_dynamically_generated(manager_info):
+            queryset_info = self.lookup_typeinfo(manager_info.metadata["django"]["from_queryset_manager"])
+            # TODO: We should always be able to find the `TypeInfo` we created a
+            # dynamic manager class from, right?
+            assert queryset_info is not None
+            arg_type = intersect_manager_and_queryset_model_argument_type(manager_info, queryset_info)
+        else:
+            arg_type = get_manager_model_argument_type(manager_info)
+
+        return arg_type
+
+    def get_manager_type(self, manager_name: str, manager_info: TypeInfo) -> Optional[MypyType]:
+        manager_arg_type = self.get_manager_model_argument_type(manager_info)
+        if (
+            isinstance(manager_arg_type, Instance)
+            and (
+                manager_arg_type.type.fullname == self.model_classdef.info.fullname
+                # We accept overwriting declared upper bound of model argument type
+                or manager_arg_type.type.fullname == fullnames.MODEL_CLASS_FULLNAME
+            )
+        ) or (
+            isinstance(manager_arg_type, AnyType) and manager_arg_type.type_of_any == TypeOfAny.from_omitted_generics
+        ):
+            return Instance(manager_info, [Instance(self.model_classdef.info, [])])
+        elif isinstance(manager_arg_type, PlaceholderType):
+            # Arg is expected to become a type, at some point, but for now we can't
+            # process it.
+            return None
+        else:
+            self.api.fail(
+                "Manager model argument does not match model it is assigned to",
+                (
+                    self.model_classdef.info.names[manager_name].node or self.model_classdef
+                    if manager_name in self.model_classdef.info.names
+                    else self.model_classdef
+                ),
+                code=MODEL_ARG_MISMATCH,
+            )
+            return Instance(manager_info, [AnyType(TypeOfAny.from_error)])
+
     def run_with_model_cls(self, model_cls: Type[Model]) -> None:
         manager_info: Optional[TypeInfo]
 
@@ -254,10 +306,14 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
                 if manager_info is None:
                     continue
 
-            is_dynamically_generated = manager_info.metadata.get("django", {}).get("from_queryset_manager") is not None
-            if manager_name not in self.model_classdef.info.names or is_dynamically_generated:
-                manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])])
-                self.add_new_node_to_model_class(manager_name, manager_type)
+            if manager_name not in self.model_classdef.info.names or self.manager_is_dynamically_generated(
+                manager_info
+            ):
+                manager_type = self.get_manager_type(manager_name, manager_info)
+                if manager_type is not None:
+                    self.add_new_node_to_model_class(manager_name, manager_type)
+                else:
+                    incomplete_manager_defs.add(manager_name)
             elif self.has_any_parametrized_manager_as_base(manager_info):
                 # Ending up here could for instance be due to having a custom _Manager_
                 # that is not built from a custom QuerySet. Another example is a
@@ -271,6 +327,13 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
                     continue
 
                 self.add_new_node_to_model_class(manager_name, custom_manager_type)
+            else:
+                manager_type = self.get_manager_type(manager_name, manager_info)
+                if manager_type is not None:
+                    # TODO: Can we avoid overwriting? Since `manager_name` is in type info names
+                    self.add_new_node_to_model_class(manager_name, manager_type)
+                else:
+                    incomplete_manager_defs.add(manager_name)
 
         if incomplete_manager_defs and not self.api.final_iteration:
             #  Unless we're on the final round, see if another round could figure out all manager types

tests/typecheck/managers/querysets/test_from_queryset.yml

@@ -584,16 +584,18 @@
 
               StrCallable = BaseManager.from_queryset(ModelQuerySet, class_name=str(1))
 
--   case: yields_error_when_mismatching_manager_and_queryset_model_args
+-   case: emits_error_when_mismatching_manager_and_queryset_model_args
     main: |
         from myapp.models import FromBaseManager, FromManager, MyModel
         reveal_type(FromBaseManager)  # N: Revealed type is "Type[myapp.models.BaseModelManagerFromModelQuerySet[Any]]"
         reveal_type(MyModel.base)  # N: Revealed type is "myapp.models.BaseModelManagerFromModelQuerySet[Any]"
         reveal_type(FromManager)  # N: Revealed type is "Type[myapp.models.ManagerModelManagerFromModelQuerySet[Any]]"
         reveal_type(MyModel.manager)  # N: Revealed type is "myapp.models.ManagerModelManagerFromModelQuerySet[Any]"
     out: |
-        myapp/models:13: error: Manager model argument not matching queryset model argument
-        myapp/models:14: error: Manager model argument not matching queryset model argument
+        myapp/models:13: error: Manager model argument does not match queryset model argument
+        myapp/models:14: error: Manager model argument does not match queryset model argument
+        myapp/models:16: error: Manager model argument does not match model it is assigned to
+        myapp/models:17: error: Manager model argument does not match model it is assigned to
     installed_apps:
         - myapp
     files:

tests/typecheck/managers/test_managers.yml

@@ -521,3 +521,89 @@
         myapp/models:47: note: Revealed type is "django.db.models.manager.BaseManager[myapp.models.InvisibleUnresolvable]"
         myapp/models:49: note: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.Booking]"
         myapp/models:50: note: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.Booking]"
+
+-   case: emits_error_when_manager_model_arg_does_not_match_model
+    main: |
+        from myapp.models import MyModel
+        reveal_type(MyModel.from_manager)  # N: Revealed type is "myapp.models.FromManager[Any]"
+        reveal_type(MyModel.from_base_manager)  # N: Revealed type is "myapp.models.FromBaseManager[Any]"
+    out: |
+        myapp/models:7: error: Manager model argument does not match model it is assigned to
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                from django.db.models.manager import BaseManager
+
+                class FromManager(models.Manager["OtherModel"]):
+                    ...
+
+                class FromBaseManager(BaseManager["OtherModel"]):
+                    ...
+
+                class MyModel(models.Model):
+                    from_manager = FromManager()
+                    from_base_manager = FromBaseManager()
+
+                class OtherModel(models.Model):
+                    ...
+
+-   case: emits_error_when_imported_manager_model_arg_does_not_match_model
+    main: |
+        from myapp.models import MyModel
+        reveal_type(MyModel.from_manager)  # N: Revealed type is "myapp.managers.FromManager[Any]"
+        reveal_type(MyModel.from_base_manager)  # N: Revealed type is "myapp.managers.FromBaseManager[Any]"
+    out: |
+        myapp/models:5: error: Manager model argument does not match model it is assigned to
+        myapp/models:6: error: Manager model argument does not match model it is assigned to
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                from .managers import FromManager, FromBaseManager
+
+                class MyModel(models.Model):
+                    from_manager = FromManager()
+                    from_base_manager = FromBaseManager()
+
+                class OtherModel(models.Model):
+                    ...
+        -   path: myapp/managers.py
+            content: |
+                from typing import TYPE_CHECKING
+                from django.db import models
+                from django.db.models.manager import BaseManager
+
+                if TYPE_CHECKING:
+                    from .models import OtherModel
+
+                class FromManager(models.Manager["OtherModel"]):
+                    ...
+
+                class FromBaseManager(BaseManager["OtherModel"]):
+                    ...
+
+-   case: populates_manager_model_arg_for_django_builtin_managers
+    main: |
+        from myapp.models import MyModel
+        reveal_type(MyModel.from_manager)  # N: Revealed type is "django.db.models.manager.Manager[myapp.models.MyModel]"
+        reveal_type(MyModel.from_base_manager)  # N: Revealed type is "django.db.models.manager.BaseManager[myapp.models.MyModel]"
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                from django.db.models.manager import BaseManager
+
+                class MyModel(models.Model):
+                    from_manager = models.Manager()
+                    # TODO: Why do get complaint that we "Need type annotation ..."?
+                    from_base_manager: BaseManager = BaseManager()