From 76e3fdb4797c1a8e764ba1e8cf767d896d5f32aa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sigurd=20Lj=C3=B8dal?=
 <544451+ljodal@users.noreply.github.com>
Date: Thu, 31 Mar 2022 22:36:22 +0200
Subject: [PATCH] Fix missing related managers on some models (#902)

* Fix missing related managers on some models

I was seeing an issue where some related managers were missing from some
models. Traced the issue down to this line, where it appears that if we
hit a relation with a non-default(?) reverse manager the iteration
stopped. I _think_ this is supposed to be a continue statement instead.
It appears to work in the project I'm working in at least.

* Add test case

* Add test case
---
 mypy_django_plugin/transformers/models.py     |  2 +-
 .../typecheck/models/test_related_fields.yml  | 89 +++++++++++++++++++
 2 files changed, 90 insertions(+), 1 deletion(-)
 create mode 100644 tests/typecheck/models/test_related_fields.yml

diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py
index 52ef3ca..7b5d705 100644
--- a/mypy_django_plugin/transformers/models.py
+++ b/mypy_django_plugin/transformers/models.py
@@ -363,7 +363,7 @@ class AddRelatedManagers(ModelClassInitializer):
                     self.add_new_node_to_model_class(
                         attname, Instance(default_reverse_manager_info, []), no_serialize=True
                     )
-                    return
+                    continue
 
                 # The reverse manager we're looking for doesn't exist. So we create it.
                 # The (default) reverse manager type is built from a RelatedManager and the default manager on the related model
diff --git a/tests/typecheck/models/test_related_fields.yml b/tests/typecheck/models/test_related_fields.yml
new file mode 100644
index 0000000..b9a8019
--- /dev/null
+++ b/tests/typecheck/models/test_related_fields.yml
@@ -0,0 +1,89 @@
+-   case: test_related_name_custom_manager
+    main: |
+      from app1.models import Model1
+      from app2.models import Model2
+
+      reveal_type(Model1().test)  # N: Revealed type is "app3.models.Model3_RelatedManager"
+      reveal_type(Model2().test)  # N: Revealed type is "app3.models.Model3_RelatedManager"
+      reveal_type(Model1().test2)  # N: Revealed type is "app3.models.Model4_RelatedManager"
+      reveal_type(Model2().test2)  # N: Revealed type is "app3.models.Model4_RelatedManager"
+    installed_apps:
+      - base
+      - users
+      - app1
+      - app2
+      - app3
+    files:
+      - path: base/__init__.py
+      - path: base/models.py
+        content: |
+          from django.db import models
+          class OwnedModel(models.Model):
+              owner = models.ForeignKey("users.User", on_delete=models.CASCADE)
+
+              class Meta:
+                  abstract = True
+
+      - path: users/__init__.py
+      - path: users/models.py
+        content: |
+          from django.db import models
+          class User(models.Model):
+              pass
+
+      - path: app1/__init__.py
+      - path: app1/models.py
+        content: |
+          from django.db import models
+          from base.models import OwnedModel
+          from typing import TYPE_CHECKING
+          if TYPE_CHECKING:
+              from app3.models import Model3
+
+          class CustomQuerySet(models.QuerySet):
+              pass
+
+          Model1Manager = models.Manager.from_queryset(CustomQuerySet)
+          class Model1(OwnedModel):
+              objects = Model1Manager()
+
+      - path: app2/__init__.py
+      - path: app2/models.py
+        content: |
+          from django.db import models
+          from base.models import OwnedModel
+          from typing import TYPE_CHECKING
+          if TYPE_CHECKING:
+              from app3.models import Model3
+
+          class CustomQuerySet(models.QuerySet):
+              pass
+
+          Model2Manager = models.Manager.from_queryset(CustomQuerySet)
+          class Model2(OwnedModel):
+              objects = Model2Manager()
+
+      - path: app3/__init__.py
+      - path: app3/models.py
+        content: |
+          from django.db import models
+          from app1.models import Model1
+          from app2.models import Model2
+          from base.models import OwnedModel
+
+          class CustomQuerySet(models.QuerySet):
+              pass
+
+          Model3Manager = models.Manager.from_queryset(CustomQuerySet)
+          class Model3(OwnedModel):
+              a = models.ForeignKey(Model1, related_name="test", on_delete=models.CASCADE)
+              b = models.ForeignKey(Model2, related_name="test", on_delete=models.CASCADE)
+
+              objects = Model3Manager()
+
+          Model4Manager = models.Manager.from_queryset(CustomQuerySet)
+          class Model4(OwnedModel):
+              a = models.ForeignKey(Model1, related_name="test2", on_delete=models.CASCADE)
+              b = models.ForeignKey(Model2, related_name="test2", on_delete=models.CASCADE)
+
+              objects = Model4Manager()