From 2a6f4647f083b2a9cd728b7e5698dc76c59b5449 Mon Sep 17 00:00:00 2001
From: Petter Friberg <petter_friberg@hotmail.com>
Date: Thu, 30 Jun 2022 09:25:29 +0200
Subject: [PATCH] Populate model argument for dynamically created managers
 (#1033)

* Populate model for dynamically created managers

* fixup! Populate model for dynamically created managers
---
 mypy_django_plugin/transformers/models.py     | 21 ++++----
 .../managers/querysets/test_from_queryset.yml | 49 +++++++++++++++++++
 2 files changed, 59 insertions(+), 11 deletions(-)

diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py
index 17a227242..6b8e50276 100644
--- a/mypy_django_plugin/transformers/models.py
+++ b/mypy_django_plugin/transformers/models.py
@@ -243,26 +243,25 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
             except helpers.IncompleteDefnException as exc:
                 # Check if manager is a generated (dynamic class) manager
                 base_manager_fullname = helpers.get_class_fullname(manager.__class__.__bases__[0])
-                if manager_fullname not in self.get_generated_manager_mappings(base_manager_fullname):
+                generated_managers = self.get_generated_manager_mappings(base_manager_fullname)
+                if manager_fullname not in generated_managers:
                     # Manager doesn't appear to be generated. Track that we encountered an
                     # incomplete definition and skip
                     incomplete_manager_defs.add(manager_name)
-                continue
+                    continue
+
+                manager_info = self.lookup_typeinfo(generated_managers[manager_fullname])
+                if manager_info is None:
+                    continue
 
-            if manager_name not in self.model_classdef.info.names:
+            is_dynamically_generated = manager_info.metadata.get("django", {}).get("from_queryset_manager") is not None
+            if manager_name not in self.model_classdef.info.names or is_dynamically_generated:
                 manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])])
                 self.add_new_node_to_model_class(manager_name, manager_type)
-            else:
+            elif self.has_any_parametrized_manager_as_base(manager_info):
                 # Ending up here could for instance be due to having a custom _Manager_
                 # that is not built from a custom QuerySet. Another example is a
                 # related manager.
-                # Don't interfere with dynamically generated manager classes
-                is_dynamically_generated = "django" in manager_info.metadata and manager_info.metadata["django"].get(
-                    "from_queryset_manager"
-                )
-                if not self.has_any_parametrized_manager_as_base(manager_info) or is_dynamically_generated:
-                    continue
-
                 custom_model_manager_name = manager.model.__name__ + "_" + manager_class_name
                 try:
                     custom_manager_type = self.create_new_model_parametrized_manager(
diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml
index 142f1f25a..2a2ae5db4 100644
--- a/tests/typecheck/managers/querysets/test_from_queryset.yml
+++ b/tests/typecheck/managers/querysets/test_from_queryset.yml
@@ -71,6 +71,55 @@
                 class MyModel(models.Model):
                     objects = NewManager()
 
+-   case: from_queryset_generated_manager_imported_from_other_module
+    main: |
+        from myapp.models import MyModel
+        reveal_type(MyModel.objects)  # N: Revealed type is "myapp.querysets.NewManager[myapp.models.MyModel]"
+        reveal_type(MyModel.objects.get())  # N: Revealed type is "myapp.models.MyModel"
+        reveal_type(MyModel.objects.queryset_method())  # N: Revealed type is "myapp.querysets.ModelQuerySet"
+        reveal_type(MyModel.objects.queryset_method_2())  # N: Revealed type is "typing.Iterable[myapp.querysets.Custom]"
+        reveal_type(MyModel.objects.queryset_method_3())  # N: Revealed type is "builtins.str"
+        reveal_type(MyModel.objects.queryset_method_4([]))  # N: Revealed type is "None"
+        reveal_type(MyModel.objects.filter(id=1).queryset_method()) # N: Revealed type is "myapp.querysets.ModelQuerySet"
+        reveal_type(MyModel.objects.filter(id=1)) # N: Revealed type is "myapp.querysets.ModelQuerySet[myapp.models.MyModel]"
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/querysets.py
+            content: |
+                from typing import TYPE_CHECKING, Iterable, Sequence
+                from django.db import models
+                from django.db.models.manager import BaseManager
+                if TYPE_CHECKING:
+                    from .models import MyModel
+
+                class Custom:
+                    ...
+
+                class ModelQuerySet(models.QuerySet["MyModel"]):
+                    def queryset_method(self) -> "ModelQuerySet":
+                        return self.filter()
+
+                    def queryset_method_2(self) -> Iterable[Custom]:
+                        return []
+
+                    def queryset_method_3(self) -> str:
+                        return 'hello'
+
+                    def queryset_method_4(self, arg: Sequence) -> None:
+                        return None
+
+                NewManager = BaseManager.from_queryset(ModelQuerySet)
+
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                from .querysets import NewManager
+
+                class MyModel(models.Model):
+                    objects = NewManager()
+
 -   case: from_queryset_with_manager
     main: |
         from myapp.models import MyModel