diff --git a/graphene/contrib/django/converter.py b/graphene/contrib/django/converter.py index 6655940a..89f90aa4 100644 --- a/graphene/contrib/django/converter.py +++ b/graphene/contrib/django/converter.py @@ -1,12 +1,22 @@ from django.db import models from ...core.types.scalars import ID, Boolean, Float, Int, String +from ...core.classtypes.enum import Enum from .compat import RelatedObject, UUIDField from .utils import get_related_model, import_single_dispatch singledispatch = import_single_dispatch() +def convert_django_field_with_choices(field): + choices = getattr(field, 'choices', None) + if choices: + meta = field.model._meta + name = '{}_{}_{}'.format(meta.app_label, meta.object_name, field.name) + return Enum(name.upper(), choices, description=field.help_text) + return convert_django_field(field) + + @singledispatch def convert_django_field(field): raise Exception( diff --git a/graphene/contrib/django/tests/test_converter.py b/graphene/contrib/django/tests/test_converter.py index 59f3aa29..1582b1ba 100644 --- a/graphene/contrib/django/tests/test_converter.py +++ b/graphene/contrib/django/tests/test_converter.py @@ -2,7 +2,8 @@ from django.db import models from py.test import raises import graphene -from graphene.contrib.django.converter import convert_django_field +from graphene.contrib.django.converter import ( + convert_django_field, convert_django_field_with_choices) from graphene.contrib.django.fields import (ConnectionOrListField, DjangoModelField) @@ -86,6 +87,26 @@ def test_should_nullboolean_convert_boolean(): assert field.required is False +def test_field_with_choices_convert_enum(): + field = models.CharField(help_text='Language', choices=( + ('es', 'Spanish'), + ('en', 'English') + )) + + class TranslatedModel(models.Model): + language = field + + class Meta: + app_label = 'test' + + graphene_type = convert_django_field_with_choices(field) + assert issubclass(graphene_type, graphene.Enum) + assert graphene_type._meta.type_name == 'TEST_TRANSLATEDMODEL_LANGUAGE' + assert graphene_type._meta.description == 'Language' + assert graphene_type.__enum__.__members__['es'].value == 'Spanish' + assert graphene_type.__enum__.__members__['en'].value == 'English' + + def test_should_float_convert_float(): assert_conversion(models.FloatField, graphene.Float) diff --git a/graphene/contrib/django/types.py b/graphene/contrib/django/types.py index 0c3bf69f..53c61f4b 100644 --- a/graphene/contrib/django/types.py +++ b/graphene/contrib/django/types.py @@ -5,7 +5,7 @@ from django.db import models from ...core.classtypes.objecttype import ObjectType, ObjectTypeMeta from ...relay.types import Connection, Node, NodeMeta -from .converter import convert_django_field +from .converter import convert_django_field_with_choices from .options import DjangoOptions from .utils import get_reverse_fields @@ -29,7 +29,7 @@ class DjangoObjectTypeMeta(ObjectTypeMeta): # We skip this field if we specify only_fields and is not # in there. Or when we exclude this field in exclude_fields continue - converted_field = convert_django_field(field) + converted_field = convert_django_field_with_choices(field) cls.add_to_class(field.name, converted_field) def construct(cls, *args, **kwargs):