From 23ad65033bc30d371061d29425ca509a919fde9e Mon Sep 17 00:00:00 2001
From: Maxim Kurnikov <maxim.kurnikov@gmail.com>
Date: Wed, 18 Sep 2019 01:40:41 +0300
Subject: [PATCH] add custom Field processing for mixins used in the Model
 subclasses (#167)

---
 mypy_django_plugin/lib/helpers.py         |  9 +++++++++
 mypy_django_plugin/transformers/fields.py |  4 +++-
 mypy_django_plugin/transformers/models.py | 15 ++++++++++++++-
 test-data/typecheck/fields/test_base.yml  | 16 ++++++++++++++++
 4 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py
index 9e7b7a5..e2e6a50 100644
--- a/mypy_django_plugin/lib/helpers.py
+++ b/mypy_django_plugin/lib/helpers.py
@@ -11,6 +11,8 @@ from mypy.plugin import (
 )
 from mypy.types import AnyType, Instance, NoneTyp, TupleType, Type as MypyType, TypeOfAny, TypedDictType, UnionType
 
+from mypy_django_plugin.lib import fullnames
+
 if TYPE_CHECKING:
     from mypy_django_plugin.django.context import DjangoContext
 
@@ -247,3 +249,10 @@ def get_typechecker_api(ctx: Union[AttributeContext, MethodContext, FunctionCont
     if not isinstance(ctx.api, TypeChecker):
         raise ValueError('Not a TypeChecker')
     return cast(TypeChecker, ctx.api)
+
+
+def get_all_model_mixins(api: TypeChecker) -> Set[str]:
+    basemodel_info = lookup_fully_qualified_typeinfo(api, fullnames.MODEL_CLASS_FULLNAME)
+    if basemodel_info is None:
+        return set()
+    return set(get_django_metadata(basemodel_info).get('model_mixins', dict).keys())
diff --git a/mypy_django_plugin/transformers/fields.py b/mypy_django_plugin/transformers/fields.py
index 1efcfe1..f91548f 100644
--- a/mypy_django_plugin/transformers/fields.py
+++ b/mypy_django_plugin/transformers/fields.py
@@ -116,7 +116,9 @@ def transform_into_proper_return_type(ctx: FunctionContext, django_context: Djan
     assert isinstance(default_return_type, Instance)
 
     outer_model_info = helpers.get_typechecker_api(ctx).scope.active_class()
-    if not outer_model_info or not outer_model_info.has_base(fullnames.MODEL_CLASS_FULLNAME):
+    if (outer_model_info is None
+            or not outer_model_info.has_base(fullnames.MODEL_CLASS_FULLNAME)
+            and outer_model_info.fullname() not in helpers.get_all_model_mixins(helpers.get_typechecker_api(ctx))):
         # not inside models.Model class
         return ctx.default_return_type
     assert isinstance(outer_model_info, TypeInfo)
diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py
index 64ce2b5..54c1469 100644
--- a/mypy_django_plugin/transformers/models.py
+++ b/mypy_django_plugin/transformers/models.py
@@ -212,6 +212,18 @@ class AddMetaOptionsAttribute(ModelClassInitializer):
                                              ]))
 
 
+class RecordAllModelMixins(ModelClassInitializer):
+    def run(self) -> None:
+        basemodel_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.MODEL_CLASS_FULLNAME)
+        basemodel_metadata = helpers.get_django_metadata(basemodel_info)
+        if 'model_mixins' not in basemodel_metadata:
+            basemodel_metadata['model_mixins'] = {}
+
+        for base_info in self.model_classdef.info.mro[1:]:
+            if base_info.fullname() != 'builtins.object':
+                basemodel_metadata['model_mixins'][base_info.fullname()] = 1
+
+
 def process_model_class(ctx: ClassDefContext,
                         django_context: DjangoContext) -> None:
     initializers = [
@@ -220,7 +232,8 @@ def process_model_class(ctx: ClassDefContext,
         AddRelatedModelsId,
         AddManagers,
         AddExtraFieldMethods,
-        AddMetaOptionsAttribute
+        AddMetaOptionsAttribute,
+        RecordAllModelMixins,
     ]
     for initializer_cls in initializers:
         try:
diff --git a/test-data/typecheck/fields/test_base.yml b/test-data/typecheck/fields/test_base.yml
index 16d98a8..9470eb6 100644
--- a/test-data/typecheck/fields/test_base.yml
+++ b/test-data/typecheck/fields/test_base.yml
@@ -132,3 +132,19 @@
             myfield: models.IntegerField[int, int]
         reveal_type(MyClass.myfield)  # N: Revealed type is 'django.db.models.fields.IntegerField[builtins.int, builtins.int]'
         reveal_type(MyClass().myfield)  # N: Revealed type is 'django.db.models.fields.IntegerField[builtins.int, builtins.int]'
+
+-   case: fields_inside_mixins_used_in_model_subclasses_resolved_as_primitives
+    main: |
+        from myapp.models import MyModel, AuthMixin
+        reveal_type(MyModel().username)  # N: Revealed type is 'builtins.str*'
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                class AuthMixin:
+                    username = models.CharField(max_length=100)
+                class MyModel(AuthMixin, models.Model):
+                    pass
\ No newline at end of file