From 356a5881e75b5bd28dbcdd52c45744fb7ef95551 Mon Sep 17 00:00:00 2001
From: Maxim Kurnikov <maxim.kurnikov@gmail.com>
Date: Sat, 4 Jan 2020 17:56:05 +0300
Subject: [PATCH] allow manager classes nested inside model classes

---
 mypy_django_plugin/lib/helpers.py             | 27 ++++++++++++++-----
 mypy_django_plugin/transformers/managers.py   |  2 +-
 scripts/enabled_test_modules.py               |  3 +++
 .../typecheck/managers/test_managers.yml      | 20 ++++++++++++++
 4 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py
index 5b7991a..6740fe3 100644
--- a/mypy_django_plugin/lib/helpers.py
+++ b/mypy_django_plugin/lib/helpers.py
@@ -31,15 +31,30 @@ def get_django_metadata(model_info: TypeInfo) -> Dict[str, Any]:
 def lookup_fully_qualified_sym(fullname: str, all_modules: Dict[str, MypyFile]) -> Optional[SymbolTableNode]:
     if '.' not in fullname:
         return None
-    module, cls_name = fullname.rsplit('.', 1)
 
-    module_file = all_modules.get(module)
+    module_file = None
+    parts = fullname.split('.')
+    for i in range(len(parts), 0, -1):
+        possible_module_name = '.'.join(parts[:i])
+        if possible_module_name in all_modules:
+            module_file = all_modules[possible_module_name]
+            break
+
     if module_file is None:
         return None
-    sym = module_file.names.get(cls_name)
-    if sym is None:
-        return None
-    return sym
+
+    cls_name = fullname.replace(module_file.fullname, '').lstrip('.')
+    sym_table = module_file.names
+    if '.' in cls_name:
+        parent_cls_name, _, cls_name = cls_name.rpartition('.')
+        # nested class
+        for parent_cls_name in parent_cls_name.split('.'):
+            sym = sym_table.get(parent_cls_name)
+            if sym is None:
+                return None
+            sym_table = sym.node.names
+
+    return sym_table.get(cls_name)
 
 
 def lookup_fully_qualified_generic(name: str, all_modules: Dict[str, MypyFile]) -> Optional[SymbolNode]:
diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py
index 6d43e4a..57de98c 100644
--- a/mypy_django_plugin/transformers/managers.py
+++ b/mypy_django_plugin/transformers/managers.py
@@ -310,4 +310,4 @@ def instantiate_anonymous_queryset_from_as_manager(ctx: MethodContext) -> MypyTy
     assert module_name == current_module.fullname
 
     generated_manager_info = current_module.names[class_name].node
-    return fill_typevars(generated_manager_info)
+    return Instance(generated_manager_info, [])
diff --git a/scripts/enabled_test_modules.py b/scripts/enabled_test_modules.py
index b1b5345..9861e6b 100644
--- a/scripts/enabled_test_modules.py
+++ b/scripts/enabled_test_modules.py
@@ -314,6 +314,9 @@ IGNORED_ERRORS = {
     'model_enums': [
         "'bool' is not a valid base class",
     ],
+    'multiple_database': [
+        'Unexpected attribute "extra_arg" for model "Book"',
+    ],
     'null_queries': [
         "Cannot resolve keyword 'foo' into field"
     ],
diff --git a/test-data/typecheck/managers/test_managers.yml b/test-data/typecheck/managers/test_managers.yml
index 22dbf33..7df83d4 100644
--- a/test-data/typecheck/managers/test_managers.yml
+++ b/test-data/typecheck/managers/test_managers.yml
@@ -335,3 +335,23 @@
                     objects = MyManager()
                 class ChildUser(models.Model):
                     objects = MyManager()
+
+
+-   case: manager_defined_in_the_nested_class
+    main: |
+        from myapp.models import MyModel
+        reveal_type(MyModel.objects)  # N: Revealed type is 'myapp.models.MyModel_MyManager[myapp.models.MyModel]'
+        reveal_type(MyModel.objects.get())  # N: Revealed type is 'myapp.models.MyModel*'
+        reveal_type(MyModel.objects.mymethod())  # 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 MyModel(models.Model):
+                    class MyManager(models.Manager):
+                        def mymethod(self) -> int:
+                            pass
+                    objects = MyManager()
\ No newline at end of file