Fix missing methods on copied manager

Author: ljodal

mypy_django_plugin/lib/helpers.py

@@ -382,29 +382,25 @@ def copy_method_to_another_class(
     method_node: FuncDef,
     return_type: Optional[MypyType] = None,
     original_module_name: Optional[str] = None,
-) -> None:
+) -> bool:
     semanal_api = get_semanal_api(ctx)
     if method_node.type is None:
-        if not semanal_api.final_iteration:
-            semanal_api.defer()
-            return
-
         arguments, return_type = build_unannotated_method_args(method_node)
         add_method_to_class(
             semanal_api, ctx.cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type
         )
-        return
+        return True
 
     method_type = method_node.type
     if not isinstance(method_type, CallableType):
         if not semanal_api.final_iteration:
             semanal_api.defer()
-        return
+        return False
 
     if return_type is None:
         return_type = bind_or_analyze_type(method_type.ret_type, semanal_api, original_module_name)
     if return_type is None:
-        return
+        return False
 
     # We build the arguments from the method signature (`CallableType`), because if we were to
     # use the arguments from the method node (`FuncDef.arguments`) we're not compatible with
@@ -417,7 +413,7 @@ def copy_method_to_another_class(
     ):
         bound_arg_type = bind_or_analyze_type(arg_type, semanal_api, original_module_name)
         if bound_arg_type is None:
-            return
+            return False
         if arg_name is None and hasattr(method_node, "arguments"):
             arg_name = method_node.arguments[pos].variable.name
         arguments.append(
@@ -435,6 +431,8 @@ def copy_method_to_another_class(
         semanal_api, ctx.cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type
     )
 
+    return True
+
 
 def add_new_manager_base(api: SemanticAnalyzerPluginInterface, fullname: str) -> None:
     sym = api.lookup_fully_qualified_or_none(fullnames.MANAGER_CLASS_FULLNAME)

mypy_django_plugin/transformers/models.py

@@ -24,6 +24,7 @@
 from mypy.types import AnyType, Instance
 from mypy.types import Type as MypyType
 from mypy.types import TypedDictType, TypeOfAny
+from typing_extensions import reveal_type
 
 from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.errorcodes import MANAGER_MISSING
@@ -211,12 +212,12 @@ def create_new_model_parametrized_manager(self, name: str, base_manager_info: Ty
         bases = []
         for original_base in base_manager_info.bases:
             if self.is_any_parametrized_manager(original_base):
-                if original_base.type is None:
-                    raise helpers.IncompleteDefnException()
-
                 original_base = helpers.reparametrize_instance(original_base, [Instance(self.model_classdef.info, [])])
             bases.append(original_base)
 
+        # TODO: This adds the manager to the module, even if we end up
+        # deferring. That can be avoided by not adding it to the module first,
+        # but rather waiting until we know we won't defer
         new_manager_info = self.add_new_class_for_current_module(name, bases)
         # copy fields to a new manager
         new_cls_def_context = ClassDefContext(cls=new_manager_info.defn, reason=self.ctx.reason, api=self.api)
@@ -225,13 +226,15 @@ def create_new_model_parametrized_manager(self, name: str, base_manager_info: Ty
         for name, sym in base_manager_info.names.items():
             # replace self type with new class, if copying method
             if isinstance(sym.node, FuncDef):
-                helpers.copy_method_to_another_class(
+                copied_method = helpers.copy_method_to_another_class(
                     new_cls_def_context,
                     self_type=custom_manager_type,
                     new_method_name=name,
                     method_node=sym.node,
                     original_module_name=base_manager_info.module_name,
                 )
+                if not copied_method and not self.api.final_iteration:
+                    raise helpers.IncompleteDefnException()
                 continue
 
             new_sym = sym.copy()

tests/typecheck/managers/test_managers.yml

@@ -364,7 +364,7 @@
 -   case: custom_manager_annotate_method_before_type_declaration
     main: |
         from myapp.models import ModelA, ModelB, ManagerA
-        reveal_type(ModelA.objects)  # N: Revealed type is "myapp.models.ModelA_ManagerA[myapp.models.ModelA]"
+        reveal_type(ModelA.objects)  # N: Revealed type is "myapp.models.ModelA_ManagerA1[myapp.models.ModelA]"
         reveal_type(ModelA.objects.do_something)  # N: Revealed type is "def (other_obj: myapp.models.ModelB) -> builtins.str"
     installed_apps:
         - myapp