add support for _meta.get_field() typechecking

mypy_django_plugin/lib/fullnames.py

@@ -35,3 +35,4 @@ RELATED_FIELDS_CLASSES = {
 }
 
 MIGRATION_CLASS_FULLNAME = 'django.db.migrations.migration.Migration'
+OPTIONS_CLASS_FULLNAME = 'django.db.models.options.Options'

mypy_django_plugin/main.py

@@ -11,7 +11,7 @@ from mypy.types import Type as MypyType
 
 from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.lib import fullnames, helpers
-from mypy_django_plugin.transformers import fields, forms, init_create, querysets, settings
+from mypy_django_plugin.transformers import fields, forms, init_create, querysets, settings, meta
 from mypy_django_plugin.transformers.models import process_model_class
 
 
@@ -168,15 +168,20 @@ class NewSemanalDjangoPlugin(Plugin):
                 return forms.extract_proper_type_for_get_form
 
         if method_name == 'values':
-            model_info = self._get_typeinfo_or_none(class_fullname)
-            if model_info and model_info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
+            info = self._get_typeinfo_or_none(class_fullname)
+            if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
                 return partial(querysets.extract_proper_type_queryset_values, django_context=self.django_context)
 
         if method_name == 'values_list':
-            model_info = self._get_typeinfo_or_none(class_fullname)
-            if model_info and model_info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
+            info = self._get_typeinfo_or_none(class_fullname)
+            if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
                 return partial(querysets.extract_proper_type_queryset_values_list, django_context=self.django_context)
 
+        if method_name == 'get_field':
+            info = self._get_typeinfo_or_none(class_fullname)
+            if info and info.has_base(fullnames.OPTIONS_CLASS_FULLNAME):
+                return partial(meta.return_proper_field_type_from_get_field, django_context=self.django_context)
+
         manager_classes = self._get_current_manager_bases()
         if class_fullname in manager_classes and method_name == 'create':
             return partial(init_create.redefine_and_typecheck_model_create, django_context=self.django_context)

mypy_django_plugin/transformers/meta.py

@@ -0,0 +1,38 @@
+from django.core.exceptions import FieldDoesNotExist
+from mypy.plugin import MethodContext
+from mypy.types import AnyType, Type as MypyType, TypeOfAny, Instance
+
+from mypy_django_plugin.django.context import DjangoContext
+from mypy_django_plugin.lib import fullnames, helpers
+
+
+def _get_field_instance(ctx: MethodContext, field_fullname: str) -> MypyType:
+    field_info = helpers.lookup_fully_qualified_typeinfo(ctx.api, field_fullname)
+    return Instance(field_info, [AnyType(TypeOfAny.explicit), AnyType(TypeOfAny.explicit)])
+
+
+def return_proper_field_type_from_get_field(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
+    model_type = ctx.type.args[0]
+    if not isinstance(model_type, Instance):
+        return _get_field_instance(ctx, fullnames.FIELD_FULLNAME)
+
+    model_cls = django_context.get_model_class_by_fullname(model_type.type.fullname())
+    if model_cls is None:
+        return _get_field_instance(ctx, fullnames.FIELD_FULLNAME)
+
+    field_name_expr = helpers.get_call_argument_by_name(ctx, 'field_name')
+    if field_name_expr is None:
+        return _get_field_instance(ctx, fullnames.FIELD_FULLNAME)
+
+    field_name = helpers.resolve_string_attribute_value(field_name_expr, ctx, django_context)
+    if field_name is None:
+        return _get_field_instance(ctx, fullnames.FIELD_FULLNAME)
+
+    try:
+        field = model_cls._meta.get_field(field_name)
+    except FieldDoesNotExist as exc:
+        ctx.api.fail(exc.args[0], ctx.context)
+        return AnyType(TypeOfAny.from_error)
+
+    field_fullname = helpers.get_class_fullname(field.__class__)
+    return _get_field_instance(ctx, field_fullname)

mypy_django_plugin/transformers/models.py

@@ -2,15 +2,15 @@ from abc import ABCMeta, abstractmethod
 from abc import ABCMeta, abstractmethod
 from typing import cast
 
-from django.db.models.fields import DateTimeField, DateField
 from django.db.models.fields.related import ForeignKey
 from django.db.models.fields.reverse_related import ManyToManyRel, ManyToOneRel, OneToOneRel
 from mypy.newsemanal.semanal import NewSemanticAnalyzer
-from mypy.nodes import MDEF, SymbolTableNode, TypeInfo, Var, Argument, ARG_NAMED_OPT, ARG_STAR2
+from mypy.nodes import ARG_STAR2, Argument, MDEF, SymbolTableNode, TypeInfo, Var
 from mypy.plugin import ClassDefContext
 from mypy.plugins import common
-from mypy.types import Instance, TypeOfAny, AnyType
+from mypy.types import AnyType, Instance, TypeOfAny
 
+from django.db.models.fields import DateField, DateTimeField
 from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.lib import fullnames, helpers
 from mypy_django_plugin.transformers import fields
@@ -177,6 +177,15 @@ class AddExtraFieldMethods(ModelClassInitializer):
                                   return_type=return_type)
 
 
+class AddMetaOptionsAttribute(ModelClassInitializer):
+    def run(self):
+        if '_meta' not in self.model_classdef.info.names:
+            options_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.OPTIONS_CLASS_FULLNAME)
+            self.add_new_node_to_model_class('_meta',
+                                             Instance(options_info, [
+                                                 Instance(self.model_classdef.info, [])
+                                             ]))
+
 
 def process_model_class(ctx: ClassDefContext,
                         django_context: DjangoContext) -> None:
@@ -186,6 +195,7 @@ def process_model_class(ctx: ClassDefContext,
         AddRelatedModelsId,
         AddManagers,
         AddExtraFieldMethods,
+        AddMetaOptionsAttribute
     ]
     for initializer_cls in initializers:
         try: