properly change type of self for methods on custom manager classes

mypy_django_plugin/transformers/models.py

@@ -1,5 +1,5 @@
 from collections import OrderedDict
-from typing import Type
+from typing import List, Tuple, Type
 
 from django.db.models.base import Model
 from django.db.models.fields import DateField, DateTimeField
@@ -7,10 +7,11 @@ from django.db.models.fields.related import ForeignKey
 from django.db.models.fields.reverse_related import (
     ManyToManyRel, ManyToOneRel, OneToOneRel,
 )
-from mypy.nodes import ARG_STAR2, Argument, Context, TypeInfo, Var
+from mypy.nodes import ARG_STAR2, Argument, Context, FuncDef, TypeInfo, Var
 from mypy.plugin import ClassDefContext
 from mypy.plugins import common
-from mypy.types import AnyType, Instance
+from mypy.plugins.common import add_method
+from mypy.types import AnyType, CallableType, Instance
 from mypy.types import Type as MypyType
 from mypy.types import TypeOfAny
 
@@ -158,7 +159,22 @@ class AddManagers(ModelClassInitializer):
                                                                            bases=bases,
                                                                            fields=OrderedDict())
                     # copy fields to a new manager
+                    new_cls_def_context = ClassDefContext(cls=custom_manager_info.defn,
+                                                          reason=self.ctx.reason,
+                                                          api=self.api)
+                    custom_manager_type = Instance(custom_manager_info, [Instance(self.model_classdef.info, [])])
+
                     for name, sym in manager_info.names.items():
+                        # replace self type with new class, if copying method
+                        if isinstance(sym.node, FuncDef):
+                            arguments, return_type = self.prepare_new_method_arguments(sym.node)
+                            add_method(new_cls_def_context,
+                                       name,
+                                       args=arguments,
+                                       return_type=return_type,
+                                       self_type=custom_manager_type)
+                            continue
+
                         new_sym = sym.copy()
                         if isinstance(new_sym.node, Var):
                             new_var = Var(name, type=sym.type)
@@ -167,9 +183,22 @@ class AddManagers(ModelClassInitializer):
                             new_sym.node = new_var
                         custom_manager_info.names[name] = new_sym
 
-                    custom_manager_type = Instance(custom_manager_info, [Instance(self.model_classdef.info, [])])
                     self.add_new_node_to_model_class(manager_name, custom_manager_type)
 
+    def prepare_new_method_arguments(self, 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
+
 
 class AddDefaultManagerAttribute(ModelClassInitializer):
     def run_with_model_cls(self, model_cls: Type[Model]) -> None:

scripts/enabled_test_modules.py

@@ -37,7 +37,6 @@ IGNORED_ERRORS = {
         'Argument after ** must be a mapping',
         'note:',
         re.compile(r'Item "None" of "[a-zA-Z_ ,\[\]]+" has no attribute'),
-        '"Optional[List[_Record]]"',
         '"Callable[..., None]" has no attribute',
         'does not return a value',
         'has no attribute "alternatives"',
@@ -257,7 +256,6 @@ IGNORED_ERRORS = {
         'Item "Field[Any, Any]" of "Union[Field[Any, Any], ForeignObjectRel]" has no attribute',
         'Incompatible types in assignment (expression has type "Type[Person',
         'Incompatible types in assignment (expression has type "FloatModel", variable has type',
-        'No overload variant of "bytes" matches argument type "Container[int]"',
         '"ImageFile" has no attribute "was_opened"',
     ],
     'model_indexes': [
@@ -278,6 +276,7 @@ IGNORED_ERRORS = {
         'Incompatible type for "department" of "Worker"',
         '"PickledModel" has no attribute',
         '"Department" has no attribute "evaluate"',
+        'Unsupported target for indexed assignment',
     ],
     'model_formsets_regress': [
         'Incompatible types in assignment (expression has type "int", target has type "str")',
@@ -352,7 +351,8 @@ IGNORED_ERRORS = {
         'Argument 1 to "append" of "list" has incompatible type "Tuple[Any, Any, Optional[Any], Any]";'
     ],
     'sites_framework': [
-        'expression has type "CurrentSiteManager[CustomArticle]", base class "AbstractArticle"'
+        'expression has type "CurrentSiteManager[CustomArticle]", base class "AbstractArticle"',
+        "Name 'Optional' is not defined",
     ],
     'syndication_tests': [
         'List or tuple expected as variable arguments'

test-data/typecheck/managers/test_managers.yml

@@ -307,9 +307,18 @@
 -   case: custom_manager_returns_proper_model_types
     main: |
         from myapp.models import User
-        reveal_type(User.objects.get())  # N: Revealed type is 'myapp.models.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.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.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'
     installed_apps:
         - myapp
     files:
@@ -318,7 +327,11 @@
             content: |
                 from django.db import models
                 class MyManager(models.Manager):
-                    def get_instance(self: "models.Manager[User]") -> int:
+                    def get_instance(self) -> int:
+                        pass
+                    def get_instance_untyped(self, name):
                         pass
                 class User(models.Model):
                     objects = MyManager()
+                class ChildUser(models.Model):
+                    objects = MyManager()