From a818a25a00754467974c91bdc5db814d4c03490d Mon Sep 17 00:00:00 2001 From: Jason Kraus Date: Thu, 18 Oct 2018 11:00:11 -0700 Subject: [PATCH] convert DRF ChoiceField to Enum, also impacts FilePathField --- graphene_django/converter.py | 27 ++++++++++++------- .../rest_framework/serializer_converter.py | 9 ++++++- .../tests/test_field_converter.py | 11 +++++--- 3 files changed, 33 insertions(+), 14 deletions(-) diff --git a/graphene_django/converter.py b/graphene_django/converter.py index c40313d..0889966 100644 --- a/graphene_django/converter.py +++ b/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) diff --git a/graphene_django/rest_framework/serializer_converter.py b/graphene_django/rest_framework/serializer_converter.py index 9f8e516..3183cd9 100644 --- a/graphene_django/rest_framework/serializer_converter.py +++ b/graphene_django/rest_framework/serializer_converter.py @@ -4,6 +4,7 @@ from rest_framework import serializers 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) diff --git a/graphene_django/rest_framework/tests/test_field_converter.py b/graphene_django/rest_framework/tests/test_field_converter.py index 6fa4ca8..2e76e32 100644 --- a/graphene_django/rest_framework/tests/test_field_converter.py +++ b/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():