convert DRF ChoiceField to Enum, also impacts FilePathField

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):
     if registry is not None:
         converted = registry.get_converted_field(field)
@@ -61,16 +77,7 @@ def convert_django_field_with_choices(field, registry=None):
     if choices:
         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)
         converted = enum(description=field.help_text, required=not field.null)
     else:
         converted = convert_django_field(field, registry)

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
 
@@ -121,7 +122,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)
 
 
@@ -138,3 +138,10 @@ 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)
+
+
+@get_graphene_type_from_serializer_field.register(serializers.ChoiceField)
+def convert_serializer_field_to_list_of_string(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,13 @@ 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 +179,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():