Make related manager inherit from objects of related model (#278)

* related manager inherits from objects of related model

* fix test typechecking

* lint

Author: mkurnikov

mypy_django_plugin/transformers/models.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Type, cast
+from typing import Dict, List, Optional, Type, cast
 
 from django.db.models.base import Model
 from django.db.models.fields import DateField, DateTimeField
@@ -58,6 +58,12 @@ def add_new_node_to_model_class(self, name: str, typ: MypyType) -> None:
                                      name=name,
                                      sym_type=typ)
 
+    def add_new_class_for_current_module(self, name: str, bases: List[Instance]) -> TypeInfo:
+        current_module = self.api.modules[self.model_classdef.info.module_name]
+        new_class_info = helpers.add_new_class_for_module(current_module,
+                                                          name=name, bases=bases)
+        return new_class_info
+
     def run(self) -> None:
         model_cls = self.django_context.get_model_class_by_fullname(self.model_classdef.fullname)
         if model_cls is None:
@@ -164,14 +170,12 @@ def create_new_model_parametrized_manager(self, name: str, base_manager_info: Ty
                                                                [Instance(self.model_classdef.info, [])])
             bases.append(original_base)
 
-        current_module = self.api.modules[self.model_classdef.info.module_name]
-        custom_manager_info = helpers.add_new_class_for_module(current_module,
-                                                               name=name, bases=bases)
+        new_manager_info = self.add_new_class_for_current_module(name, bases)
         # copy fields to a new manager
-        new_cls_def_context = ClassDefContext(cls=custom_manager_info.defn,
+        new_cls_def_context = ClassDefContext(cls=new_manager_info.defn,
                                               reason=self.ctx.reason,
                                               api=self.api)
-        custom_manager_type = Instance(custom_manager_info, [Instance(self.model_classdef.info, [])])
+        custom_manager_type = Instance(new_manager_info, [Instance(self.model_classdef.info, [])])
 
         for name, sym in base_manager_info.names.items():
             # replace self type with new class, if copying method
@@ -185,10 +189,10 @@ def create_new_model_parametrized_manager(self, name: str, base_manager_info: Ty
             new_sym = sym.copy()
             if isinstance(new_sym.node, Var):
                 new_var = Var(name, type=sym.type)
-                new_var.info = custom_manager_info
-                new_var._fullname = custom_manager_info.fullname + '.' + name
+                new_var.info = new_manager_info
+                new_var._fullname = new_manager_info.fullname + '.' + name
                 new_sym.node = new_var
-            custom_manager_info.names[name] = new_sym
+            new_manager_info.names[name] = new_sym
 
         return custom_manager_type
 
@@ -268,15 +272,30 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
 
             if isinstance(relation, (ManyToOneRel, ManyToManyRel)):
                 try:
-                    manager_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.RELATED_MANAGER_CLASS)
+                    related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.RELATED_MANAGER_CLASS)  # noqa: E501
+                    if 'objects' not in related_model_info.names:
+                        raise helpers.IncompleteDefnException()
                 except helpers.IncompleteDefnException as exc:
                     if not self.api.final_iteration:
                         raise exc
                     else:
                         continue
-                self.add_new_node_to_model_class(attname,
-                                                 Instance(manager_info, [Instance(related_model_info, [])]))
-                continue
+
+                # create new RelatedManager subclass
+                parametrized_related_manager_type = Instance(related_manager_info,
+                                                             [Instance(related_model_info, [])])
+                default_manager_type = related_model_info.names['objects'].type
+                if (default_manager_type is None
+                        or not isinstance(default_manager_type, Instance)
+                        or default_manager_type.type.fullname == fullnames.MANAGER_CLASS_FULLNAME):
+                    self.add_new_node_to_model_class(attname, parametrized_related_manager_type)
+                    continue
+
+                name = related_model_cls.__name__ + '_' + 'RelatedManager'
+                bases = [parametrized_related_manager_type, default_manager_type]
+                new_related_manager_info = self.add_new_class_for_current_module(name, bases)
+
+                self.add_new_node_to_model_class(attname, Instance(new_related_manager_info, []))
 
 
 class AddExtraFieldMethods(ModelClassInitializer):

scripts/enabled_test_modules.py

@@ -248,7 +248,7 @@
         re.compile(r'"(HttpRequest|WSGIRequest)" has no attribute'),
     ],
     'many_to_many': [
-        '(expression has type "List[Article]", variable has type "RelatedManager[Article]"',
+        '(expression has type "List[Article]", variable has type "Article_RelatedManager2',
         '"add" of "RelatedManager" has incompatible type "Article"; expected "Union[Publication, int]"',
     ],
     'many_to_one': [

test-data/typecheck/fields/test_related.yml

@@ -648,3 +648,27 @@
                         abstract = True
                 class User(AbstractUser):
                     pass
+
+
+-   case: related_manager_is_a_subclass_of_default_manager
+    main: |
+        from myapp.models import User
+        reveal_type(User().orders)  # N: Revealed type is 'myapp.models.Order_RelatedManager'
+        reveal_type(User().orders.get())  # N: Revealed type is 'myapp.models.Order*'
+        reveal_type(User().orders.manager_method())  # N: Revealed type is 'builtins.int'
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                class User(models.Model):
+                    pass
+                class OrderManager(models.Manager):
+                    def manager_method(self) -> int:
+                        pass
+                class Order(models.Model):
+                    objects = OrderManager()
+                    user = models.ForeignKey(to=User, on_delete=models.CASCADE, related_name='orders')
+