Create related managers from generated managers (#580)

mypy_django_plugin/transformers/models.py

@@ -275,6 +275,8 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
                 # 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:
+                    default_manager_type = self.try_generate_related_manager(related_model_cls, related_model_info)
                 if (
                     default_manager_type is None
                     or not isinstance(default_manager_type, Instance)
@@ -283,12 +285,32 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
                     self.add_new_node_to_model_class(attname, parametrized_related_manager_type)
                     continue
 
-                name = related_model_cls.__name__ + "_" + "RelatedManager"
+                name = model_cls.__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, []))
 
+    def get_generated_manager_mappings(self, base_manager_fullname: str) -> Dict[str, str]:
+        base_manager_info = self.lookup_typeinfo(base_manager_fullname)
+        if base_manager_info is None or "from_queryset_managers" not in base_manager_info.metadata:
+            return {}
+        return base_manager_info.metadata["from_queryset_managers"]
+
+    def try_generate_related_manager(
+        self, related_model_cls: Type[Model], related_model_info: TypeInfo
+    ) -> Optional[Instance]:
+        manager = related_model_cls._meta.managers_map["objects"]
+        base_manager_fullname = helpers.get_class_fullname(manager.__class__.__bases__[0])
+        manager_fullname = helpers.get_class_fullname(manager.__class__)
+        generated_managers = self.get_generated_manager_mappings(base_manager_fullname)
+        if manager_fullname in generated_managers:
+            real_manager_fullname = generated_managers[manager_fullname]
+            manager_info = self.lookup_typeinfo(real_manager_fullname)  # type: ignore
+            if manager_info:
+                return Instance(manager_info, [Instance(related_model_info, [])])
+        return None
+
 
 class AddExtraFieldMethods(ModelClassInitializer):
     def run_with_model_cls(self, model_cls: Type[Model]) -> None:

scripts/enabled_test_modules.py

@@ -290,7 +290,7 @@
         'Incompatible types in assignment (expression has type "HttpResponseBase", variable has type "HttpResponse")',
     ],
     "many_to_many": [
-        '(expression has type "List[Article]", variable has type "Article_RelatedManager2',
+        '(expression has type "List[Article]", variable has type "Publication_Article_RelatedManager1',
         '"add" of "RelatedManager" has incompatible type "Article"; expected "Union[Publication, int]"',
     ],
     "many_to_one": [

tests/typecheck/fields/test_related.yml

@@ -652,10 +652,18 @@
 
 -   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'
+        from myapp.models import User, Order, Product
+        reveal_type(User().orders)  # N: Revealed type is 'myapp.models.User_Order_RelatedManager1'
         reveal_type(User().orders.get())  # N: Revealed type is 'myapp.models.Order*'
         reveal_type(User().orders.manager_method())  # N: Revealed type is 'builtins.int'
+        reveal_type(Order().products)  # N: Revealed type is 'myapp.models.Order_Product_RelatedManager1'
+        reveal_type(Order().products.get())  # N: Revealed type is 'myapp.models.Product*'
+        reveal_type(Order().products.queryset_method())  # N: Revealed type is 'builtins.int'
+        # TODO: realted manager support to use the same type for all related managers
+        if 1 == 2:
+            manager = User().products
+        else:
+            manager = Order().products  # E: Incompatible types in assignment (expression has type "Order_Product_RelatedManager1", variable has type "User_Product_RelatedManager1")
     installed_apps:
         - myapp
     files:
@@ -671,6 +679,46 @@
                 class Order(models.Model):
                     objects = OrderManager()
                     user = models.ForeignKey(to=User, on_delete=models.CASCADE, related_name='orders')
+                class ProductQueryset(models.QuerySet):
+                    def queryset_method(self) -> int:
+                        pass
+                ProductManager = models.Manager.from_queryset(ProductQueryset)
+                class Product(models.Model):
+                    objects = ProductManager()
+                    order = models.ForeignKey(to=Order, on_delete=models.CASCADE, related_name='products')
+                    user = models.ForeignKey(to=User, on_delete=models.CASCADE, related_name='products')
+
+-   case: related_manager_no_conflict_from_star_import
+    main: |
+        import myapp.models
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models/__init__.py
+            content: |
+                from myapp.models.a import *
+                # make sure generated related manager from address to user doesn't have
+                # the same name with related manager from profile to user
+                from myapp.models.b import *
+        -   path: myapp/models/a.py
+            content: |
+                from django.db import models
+                class Address(models.Model):
+                    pass
+        -   path: myapp/models/b.py
+            content: |
+                from django.db import models
+                from .a import Address
+                class Profile(models.Model):
+                    pass
+                class UserQuerySet(models.QuerySet):
+                    pass
+                UserManager = models.Manager.from_queryset(UserQuerySet)
+                class User(models.Model):
+                    address = models.ForeignKey(Address, on_delete=models.CASCADE)
+                    profile = models.ForeignKey(Profile, on_delete=models.CASCADE)
+                    objects = UserManager()
 
 -   case: many_to_many_field_can_be_used_in_alias
     main: |