diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index e0cca53..300760e 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/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: diff --git a/test-data/typecheck/models/test_extra_methods.yml b/test-data/typecheck/models/test_extra_methods.yml new file mode 100644 index 0000000..a920ada --- /dev/null +++ b/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)