add support for get_FIELD_display for choices fields

mypy_django_plugin/transformers/models.py

@@ -7,6 +7,7 @@ from django.db.models.fields.reverse_related import ManyToManyRel, ManyToOneRel,
 from mypy.newsemanal.semanal import NewSemanticAnalyzer
 from mypy.nodes import MDEF, SymbolTableNode, TypeInfo, Var
 from mypy.plugin import ClassDefContext
+from mypy.plugins import common
 from mypy.types import Instance
 
 from mypy_django_plugin.django.context import DjangoContext
@@ -139,6 +140,22 @@ class AddManagers(ModelClassInitializer):
                 continue
 
 
+class AddFieldChoicesDisplayMethods(ModelClassInitializer):
+    def run(self):
+        model_cls = self.django_context.get_model_class_by_fullname(self.model_classdef.fullname)
+        if model_cls is None:
+            return
+
+        for field in self.django_context.get_model_fields(model_cls):
+            if field.choices:
+                info = self.lookup_typeinfo_or_incomplete_defn_error('builtins.str')
+                return_type = Instance(info, [])
+                common.add_method(self.ctx,
+                                  name='get_{}_display'.format(field.attname),
+                                  args=[],
+                                  return_type=return_type)
+
+
 def process_model_class(ctx: ClassDefContext,
                         django_context: DjangoContext) -> None:
     initializers = [
@@ -146,6 +163,7 @@ def process_model_class(ctx: ClassDefContext,
         AddDefaultPrimaryKey,
         AddRelatedModelsId,
         AddManagers,
+        AddFieldChoicesDisplayMethods,
     ]
     for initializer_cls in initializers:
         try:

test-data/typecheck/models/test_extra_methods.yml

@@ -0,0 +1,20 @@
+-   case: if_field_has_choices_set_model_has_get_FIELDNAME_display_method
+    main: |
+        from myapp.models import MyUser
+        user = MyUser(name='user', gender='M')
+        user.get_name_display()  # E: "MyUser" has no attribute "get_name_display"; maybe "get_gender_display"?
+        reveal_type(user.get_gender_display())  # N: Revealed type is 'builtins.str'
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                GENDER_CHOICES = (
+                    ('M', 'Male'),
+                    ('F', 'Female'),
+                )
+                class MyUser(models.Model):
+                    name = models.CharField(max_length=100)
+                    gender = models.CharField(max_length=100, choices=GENDER_CHOICES)