support for some Field -> builtin type mapping, OneToOneField, ForeignKey

mypy_django_plugin/plugins/base.py

@@ -0,0 +1,77 @@
+from typing import Callable, Optional
+
+from mypy.nodes import AssignmentStmt, CallExpr, RefExpr, StrExpr
+from mypy.plugin import Plugin, ClassDefContext
+
+from mypy_django_plugin.helpers import get_app_model
+from mypy_django_plugin.model_classes import DjangoModelsRegistry
+
+
+# fields which real type is inside to= expression
+REFERENCING_DB_FIELDS = {
+    'django.db.models.fields.related.ForeignKey',
+    'django.db.models.fields.related.OneToOneField'
+}
+
+
+def save_referred_to_model_in_metadata(rvalue: CallExpr) -> None:
+    to_arg_value = rvalue.args[rvalue.arg_names.index('to')]
+    if isinstance(to_arg_value, StrExpr):
+        referred_model_fullname = get_app_model(to_arg_value.value)
+    else:
+        referred_model_fullname = to_arg_value.fullname
+
+    rvalue.callee.node.metadata['base'] = referred_model_fullname
+
+
+class CollectModelsInformation(object):
+    def __init__(self, model_registry: DjangoModelsRegistry):
+        self.model_registry = model_registry
+
+    def __call__(self, model_definition: ClassDefContext) -> None:
+        self.model_registry.base_models.add(model_definition.cls.fullname)
+
+        for member in model_definition.cls.defs.body:
+            if isinstance(member, AssignmentStmt):
+                if len(member.lvalues) > 1:
+                    return None
+
+                arg_name = member.lvalues[0].name
+                arg_name_as_id = arg_name + '_id'
+
+                rvalue = member.rvalue
+                if isinstance(rvalue, CallExpr):
+                    if not isinstance(rvalue.callee, RefExpr):
+                        return None
+
+                    if rvalue.callee.fullname in REFERENCING_DB_FIELDS:
+                        if rvalue.callee.fullname == 'django.db.models.fields.related.ForeignKey':
+                            model_definition.cls.info.names[arg_name_as_id] = \
+                                model_definition.api.lookup_fully_qualified('builtins.int')
+
+                        if rvalue.callee.fullname == 'django.db.models.fields.related.OneToOneField':
+                            if 'related_name' in rvalue.arg_names:
+                                referred_to_model = rvalue.args[rvalue.arg_names.index('to')]
+                                related_arg_value = rvalue.args[rvalue.arg_names.index('related_name')].value
+
+                                if isinstance(referred_to_model, StrExpr):
+                                    referred_model_fullname = get_app_model(referred_to_model.value)
+                                else:
+                                    referred_model_fullname = referred_to_model.fullname
+
+                                referred_model = model_definition.api.lookup_fully_qualified_or_none(referred_model_fullname)
+                                referred_model.node.names[related_arg_value] = \
+                                    model_definition.api.lookup_fully_qualified_or_none(model_definition.cls.fullname)
+
+                        return save_referred_to_model_in_metadata(rvalue)
+
+
+class BaseDjangoModelsPlugin(Plugin):
+    model_registry = DjangoModelsRegistry()
+
+    def get_base_class_hook(self, fullname: str
+                            ) -> Optional[Callable[[ClassDefContext], None]]:
+        if fullname in self.model_registry:
+            return CollectModelsInformation(self.model_registry)
+
+        return None