From 9938378e94c21b34c4a9425fb3d93a166f034b7d Mon Sep 17 00:00:00 2001
From: Rob Percival <rob.percival@krakentechnologies.ltd>
Date: Tue, 19 Oct 2021 12:31:14 +0100
Subject: [PATCH] Make AddRelatedManagers look for "objects" on parent model
 (#730)

* Add failing test for relation to model inheriting `objects`

Fails with:
```
pytest_mypy_plugins.utils.TypecheckAssertionError: Invalid output:
Expected:
  main:2: note: Revealed type is "myapp.models.MyUser*"
  main:3: note: Revealed type is "myapp.models.MyUser*"
  <45 (diff)
  <45 (diff)
Actual:
  main:2: note: Revealed type is "myapp.models.MyUser*"
  main:3: note: Revealed type is "myapp.models.MyUser*"
  main:6: error: "MyUser" has no attribute "book_set" (diff)
  main:6: note: Revealed type is "Any"          (diff)
  main:7: error: "MyUser" has no attribute "article_set" (diff)
  main:7: note: Revealed type is "Any"          (diff)
```

* Make AddRelatedManagers look for "objects" on parent model

Previously, AddRelatedManagers would fail if a related model had inherited
its `objects` field from a parent class. This would result in missing
relation attributes. This is fixed by using `get()` instead of `names`;
the former searches the MRO for the symbol, whereas the latter only looks
for symbols declared directly on the class.
---
 mypy_django_plugin/transformers/models.py |  5 ++--
 tests/typecheck/fields/test_related.yml   | 31 +++++++++++++++++++++++
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py
index 305ade9..73ccf93 100644
--- a/mypy_django_plugin/transformers/models.py
+++ b/mypy_django_plugin/transformers/models.py
@@ -273,7 +273,8 @@ class AddRelatedManagers(ModelClassInitializer):
                     related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(
                         fullnames.RELATED_MANAGER_CLASS
                     )  # noqa: E501
-                    if "objects" not in related_model_info.names:
+                    objects = related_model_info.get("objects")
+                    if not objects:
                         raise helpers.IncompleteDefnException()
                 except helpers.IncompleteDefnException as exc:
                     if not self.api.final_iteration:
@@ -283,7 +284,7 @@ class AddRelatedManagers(ModelClassInitializer):
 
                 # 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
+                default_manager_type = objects.type
                 if default_manager_type is None:
                     default_manager_type = self.try_generate_related_manager(related_model_cls, related_model_info)
                 if (
diff --git a/tests/typecheck/fields/test_related.yml b/tests/typecheck/fields/test_related.yml
index c54b33e..04a42c7 100644
--- a/tests/typecheck/fields/test_related.yml
+++ b/tests/typecheck/fields/test_related.yml
@@ -601,6 +601,37 @@
                     pass
 
 
+-   case: test_foreign_key_from_superclass_inherits_correctly_when_also_inheriting_manager
+    main: |
+        from myapp.models import MyUser, Book, Article, LibraryEntity
+        reveal_type(Book().registered_by_user)  # N: Revealed type is "myapp.models.MyUser*"
+        reveal_type(Article().registered_by_user)  # N: Revealed type is "myapp.models.MyUser*"
+
+        user = MyUser()
+        reveal_type(user.book_set) # N: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.Book]"
+        reveal_type(user.article_set) # N: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.Article]"
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+              from django.db import models
+              class MyUser(models.Model):
+                  pass
+              class LibraryEntityQuerySet(models.QuerySet):
+                  pass
+              class LibraryEntity(models.Model):
+                  class Meta:
+                      abstract = True
+                  objects = models.Manager.from_queryset(LibraryEntityQuerySet)()
+                  registered_by_user = models.ForeignKey(MyUser, on_delete=models.CASCADE)
+              class Book(LibraryEntity):
+                  pass
+              class Article(LibraryEntity):
+                  pass
+
+
 -   case: foreign_key_relationship_for_models_with_custom_manager
     main: |
         from myapp.models import Transaction