convert DRF ChoiceField to Enum (#537)

* convert DRF ChoiceField to Enum, also impacts FilePathField

* Pep8 fixes

* DRF multiple choices field converts to list of enum

* apply black formatting

Author: zbyte64

graphene_django/converter.py

@@ -1,3 +1,4 @@
+from collections import OrderedDict
 from django.db import models
 from django.utils.encoding import force_text
 
@@ -39,6 +40,8 @@ def convert_choice_name(name):
 
 def get_choices(choices):
     converted_names = []
+    if isinstance(choices, OrderedDict):
+        choices = choices.items()
     for value, help_text in choices:
         if isinstance(help_text, (tuple, list)):
             for choice in get_choices(help_text):
@@ -52,6 +55,19 @@ def get_choices(choices):
             yield name, value, description
 
 
+def convert_choices_to_named_enum_with_descriptions(name, choices):
+    choices = list(get_choices(choices))
+    named_choices = [(c[0], c[1]) for c in choices]
+    named_choices_descriptions = {c[0]: c[2] for c in choices}
+
+    class EnumWithDescriptionsType(object):
+        @property
+        def description(self):
+            return named_choices_descriptions[self.name]
+
+    return Enum(name, list(named_choices), type=EnumWithDescriptionsType)
+
+
 def convert_django_field_with_choices(
     field, registry=None, convert_choices_to_enum=True
 ):
@@ -63,16 +79,7 @@ def convert_django_field_with_choices(
     if choices and convert_choices_to_enum:
         meta = field.model._meta
         name = to_camel_case("{}_{}".format(meta.object_name, field.name))
-        choices = list(get_choices(choices))
-        named_choices = [(c[0], c[1]) for c in choices]
-        named_choices_descriptions = {c[0]: c[2] for c in choices}
-
-        class EnumWithDescriptionsType(object):
-            @property
-            def description(self):
-                return named_choices_descriptions[self.name]
-
-        enum = Enum(name, list(named_choices), type=EnumWithDescriptionsType)
+        enum = convert_choices_to_named_enum_with_descriptions(name, choices)
         required = not (field.blank or field.null)
         converted = enum(description=field.help_text, required=required)
     else:

graphene_django/rest_framework/serializer_converter.py

@@ -4,6 +4,7 @@
 import graphene
 
 from ..registry import get_global_registry
+from ..converter import convert_choices_to_named_enum_with_descriptions
 from ..utils import import_single_dispatch
 from .types import DictType
 
@@ -130,7 +131,6 @@ def convert_serializer_field_to_time(field):
 @get_graphene_type_from_serializer_field.register(serializers.ListField)
 def convert_serializer_field_to_list(field, is_input=True):
     child_type = get_graphene_type_from_serializer_field(field.child)
-
     return (graphene.List, child_type)
 
 
@@ -145,5 +145,13 @@ def convert_serializer_field_to_jsonstring(field):
 
 
 @get_graphene_type_from_serializer_field.register(serializers.MultipleChoiceField)
-def convert_serializer_field_to_list_of_string(field):
-    return (graphene.List, graphene.String)
+def convert_serializer_field_to_list_of_enum(field):
+    child_type = convert_serializer_field_to_enum(field)
+    return (graphene.List, child_type)
+
+
+@get_graphene_type_from_serializer_field.register(serializers.ChoiceField)
+def convert_serializer_field_to_enum(field):
+    # enums require a name
+    name = field.field_name or field.source or "Choices"
+    return convert_choices_to_named_enum_with_descriptions(name, field.choices)

graphene_django/rest_framework/tests/test_field_converter.py

@@ -60,8 +60,17 @@ def test_should_url_convert_string():
     assert_conversion(serializers.URLField, graphene.String)
 
 
-def test_should_choice_convert_string():
-    assert_conversion(serializers.ChoiceField, graphene.String, choices=[])
+def test_should_choice_convert_enum():
+    field = assert_conversion(
+        serializers.ChoiceField,
+        graphene.Enum,
+        choices=[("h", "Hello"), ("w", "World")],
+        source="word",
+    )
+    assert field._meta.enum.__members__["H"].value == "h"
+    assert field._meta.enum.__members__["H"].description == "Hello"
+    assert field._meta.enum.__members__["W"].value == "w"
+    assert field._meta.enum.__members__["W"].description == "World"
 
 
 def test_should_base_field_convert_string():
@@ -174,7 +183,7 @@ def test_should_file_convert_string():
 
 
 def test_should_filepath_convert_string():
-    assert_conversion(serializers.FilePathField, graphene.String, path="/")
+    assert_conversion(serializers.FilePathField, graphene.Enum, path="/")
 
 
 def test_should_ip_convert_string():
@@ -189,9 +198,9 @@ def test_should_json_convert_jsonstring():
     assert_conversion(serializers.JSONField, graphene.types.json.JSONString)
 
 
-def test_should_multiplechoicefield_convert_to_list_of_string():
+def test_should_multiplechoicefield_convert_to_list_of_enum():
     field = assert_conversion(
         serializers.MultipleChoiceField, graphene.List, choices=[1, 2, 3]
     )
 
-    assert field.of_type == graphene.String
+    assert issubclass(field.of_type, graphene.Enum)