BaseManager.from_queryset(): properly resolve methods for QuerySet defined in another file (#282)

* BaseManager.from_queryset() from another file

* only anal_type per argument

* add resolve for return_type

* fix mypy errors

* remove leftover comment

mypy_django_plugin/lib/helpers.py

@@ -10,8 +10,8 @@ from mypy import checker
 from mypy.checker import TypeChecker
 from mypy.mro import calculate_mro
 from mypy.nodes import (
-    GDEF, MDEF, Argument, Block, ClassDef, Expression, FuncDef, MemberExpr, MypyFile, NameExpr, StrExpr, SymbolNode,
-    SymbolTable, SymbolTableNode, TypeInfo, Var,
+    GDEF, MDEF, Argument, Block, ClassDef, Expression, FuncDef, MemberExpr, MypyFile, NameExpr, PlaceholderNode,
+    StrExpr, SymbolNode, SymbolTable, SymbolTableNode, TypeInfo, Var,
 )
 from mypy.plugin import (
     AttributeContext, CheckerPluginInterface, ClassDefContext, DynamicClassDefContext, FunctionContext, MethodContext,
@@ -309,39 +309,67 @@ def add_new_sym_for_info(info: TypeInfo, *, name: str, sym_type: MypyType) -> No
                                        plugin_generated=True)
 
 
-def _prepare_new_method_arguments(node: FuncDef) -> Tuple[List[Argument], MypyType]:
-    arguments = []
-    for argument in node.arguments[1:]:
-        if argument.type_annotation is None:
-            argument.type_annotation = AnyType(TypeOfAny.unannotated)
-        arguments.append(argument)
-
-    if isinstance(node.type, CallableType):
-        return_type = node.type.ret_type
-    else:
-        return_type = AnyType(TypeOfAny.unannotated)
-
-    return arguments, return_type
+def build_unannotated_method_args(method_node: FuncDef) -> Tuple[List[Argument], MypyType]:
+    prepared_arguments = []
+    for argument in method_node.arguments[1:]:
+        argument.type_annotation = AnyType(TypeOfAny.unannotated)
+        prepared_arguments.append(argument)
+    return_type = AnyType(TypeOfAny.unannotated)
+    return prepared_arguments, return_type
 
 
 def copy_method_to_another_class(ctx: ClassDefContext, self_type: Instance,
                                  new_method_name: str, method_node: FuncDef) -> None:
-    arguments, return_type = _prepare_new_method_arguments(method_node)
-
     semanal_api = get_semanal_api(ctx)
-    for argument in arguments:
-        if argument.type_annotation is not None:
-            argument.type_annotation = semanal_api.anal_type(argument.type_annotation,
-                                                             allow_placeholder=True)
+    if method_node.type is None:
+        if not semanal_api.final_iteration:
+            semanal_api.defer()
+            return
 
-    if return_type is not None:
-        ret = semanal_api.anal_type(return_type,
-                                    allow_placeholder=True)
-        assert ret is not None
-        return_type = ret
+        arguments, return_type = build_unannotated_method_args(method_node)
+        add_method(ctx,
+                   new_method_name,
+                   args=arguments,
+                   return_type=return_type,
+                   self_type=self_type)
+        return
+
+    method_type = method_node.type
+    if not isinstance(method_type, CallableType):
+        if not semanal_api.final_iteration:
+            semanal_api.defer()
+        return
+
+    arguments = []
+    bound_return_type = semanal_api.anal_type(method_type.ret_type,
+                                              allow_placeholder=True)
+    assert bound_return_type is not None
+
+    if isinstance(bound_return_type, PlaceholderNode):
+        return
+
+    for arg_name, arg_type, original_argument in zip(method_type.arg_names[1:],
+                                                     method_type.arg_types[1:],
+                                                     method_node.arguments[1:]):
+        bound_arg_type = semanal_api.anal_type(arg_type, allow_placeholder=True)
+        assert bound_arg_type is not None
+
+        if isinstance(bound_arg_type, PlaceholderNode):
+            return
+
+        var = Var(name=original_argument.variable.name,
+                  type=arg_type)
+        var.line = original_argument.variable.line
+        var.column = original_argument.variable.column
+        argument = Argument(variable=var,
+                            type_annotation=bound_arg_type,
+                            initializer=original_argument.initializer,
+                            kind=original_argument.kind)
+        argument.set_line(original_argument)
+        arguments.append(argument)
 
     add_method(ctx,
                new_method_name,
                args=arguments,
-               return_type=return_type,
+               return_type=bound_return_type,
                self_type=self_type)

test-data/typecheck/managers/querysets/test_from_queryset.yml

@@ -146,3 +146,36 @@
                         return param
 
                 NewManager = models.Manager.from_queryset(ModelQuerySet)
+
+-   case: from_queryset_with_inherited_manager_and_typing_no_return
+    disable_cache: true
+    main: |
+        from myapp.models import MyModel
+        reveal_type(MyModel().objects)  # N: Revealed type is 'myapp.models.MyModel_NewManager[myapp.models.MyModel]'
+        reveal_type(MyModel().objects.get())  # N: Revealed type is 'myapp.models.MyModel*'
+        reveal_type(MyModel().objects.base_queryset_method)  # N: Revealed type is 'def (param: Union[builtins.int, builtins.str]) -> <nothing>'
+        reveal_type(MyModel().objects.base_queryset_method(2))  # N: Revealed type is '<nothing>'
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                from myapp.managers import NewManager
+                class MyModel(models.Model):
+                    objects = NewManager()
+        -   path: myapp/managers.py
+            content: |
+                from django.db import models
+                from myapp.base_queryset import BaseQuerySet
+                class ModelQuerySet(BaseQuerySet):
+                    pass
+                NewManager = models.Manager.from_queryset(ModelQuerySet)
+        -   path: myapp/base_queryset.py
+            content: |
+                from typing import NoReturn, Union
+                from django.db import models
+                class BaseQuerySet(models.QuerySet):
+                    def base_queryset_method(self, param: Union[int, str]) -> NoReturn:
+                        raise ValueError

test-data/typecheck/managers/test_managers.yml

@@ -307,15 +307,15 @@
 -   case: custom_manager_returns_proper_model_types
     main: |
         from myapp.models import User
-        reveal_type(User.objects)  # N: Revealed type is 'myapp.models.User_MyManager[myapp.models.User]'
-        reveal_type(User.objects.select_related())  # N: Revealed type is 'myapp.models.User_MyManager[myapp.models.User]'
+        reveal_type(User.objects)  # N: Revealed type is 'myapp.models.User_MyManager2[myapp.models.User]'
+        reveal_type(User.objects.select_related())  # N: Revealed type is 'myapp.models.User_MyManager2[myapp.models.User]'
         reveal_type(User.objects.get())  # N: Revealed type is 'myapp.models.User*'
         reveal_type(User.objects.get_instance())  # N: Revealed type is 'builtins.int'
         reveal_type(User.objects.get_instance_untyped('hello'))  # N: Revealed type is 'Any'
 
         from myapp.models import ChildUser
-        reveal_type(ChildUser.objects)  # N: Revealed type is 'myapp.models.ChildUser_MyManager[myapp.models.ChildUser]'
-        reveal_type(ChildUser.objects.select_related())  # N: Revealed type is 'myapp.models.ChildUser_MyManager[myapp.models.ChildUser]'
+        reveal_type(ChildUser.objects)  # N: Revealed type is 'myapp.models.ChildUser_MyManager2[myapp.models.ChildUser]'
+        reveal_type(ChildUser.objects.select_related())  # N: Revealed type is 'myapp.models.ChildUser_MyManager2[myapp.models.ChildUser]'
         reveal_type(ChildUser.objects.get())  # N: Revealed type is 'myapp.models.ChildUser*'
         reveal_type(ChildUser.objects.get_instance())  # N: Revealed type is 'builtins.int'
         reveal_type(ChildUser.objects.get_instance_untyped('hello'))  # N: Revealed type is 'Any'