From 59405563d7588df005f757f5386597e970dd092d Mon Sep 17 00:00:00 2001 From: Philip Tzou Date: Mon, 10 Oct 2016 11:31:38 -0700 Subject: [PATCH] Serialization function should be applied on `ChoiceType` output --- graphene_sqlalchemy/converter.py | 22 +++++++++++++++++++-- graphene_sqlalchemy/tests/test_converter.py | 14 +++++++++++-- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/graphene_sqlalchemy/converter.py b/graphene_sqlalchemy/converter.py index 05690697..7ac16a9e 100644 --- a/graphene_sqlalchemy/converter.py +++ b/graphene_sqlalchemy/converter.py @@ -10,6 +10,12 @@ from .fields import SQLAlchemyConnectionField +try: + from enum import Enum as PyEnum +except ImportError: + # ChoiceType supports built-in enum and enum34 + PyEnum = None + try: from sqlalchemy_utils import ChoiceType, JSONType, ScalarListType except ImportError: @@ -116,9 +122,21 @@ def convert_column_to_float(type, column, registry=None): @convert_sqlalchemy_type.register(ChoiceType) -def convert_column_to_enum(type, column, registry=None): +def convert_column_to_enum(type_, column, registry=None): name = '{}_{}'.format(column.table.name, column.name).upper() - return Enum(name, type.choices, description=column.doc) + if PyEnum is not None and isinstance(type_.choices, type) and issubclass(type_.choices, PyEnum): + enum_class = Enum.from_enum(type_.choices, description=column.doc) + else: + enum_class = Enum(name, type_.choices, description=column.doc) + + def serialize(self, value): + enum_value = self._value_lookup.get(value.value) + if enum_value: + return enum_value.name + return None + + enum_class.serialize = staticmethod(serialize) + return enum_class(required=not(column.nullable)) @convert_sqlalchemy_type.register(ScalarListType) diff --git a/graphene_sqlalchemy/tests/test_converter.py b/graphene_sqlalchemy/tests/test_converter.py index 0157f837..e584b75b 100644 --- a/graphene_sqlalchemy/tests/test_converter.py +++ b/graphene_sqlalchemy/tests/test_converter.py @@ -4,7 +4,7 @@ from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import composite from sqlalchemy.sql.elements import Label -from sqlalchemy_utils import ChoiceType, JSONType, ScalarListType +from sqlalchemy_utils import Choice, ChoiceType, JSONType, ScalarListType import graphene from graphene.relay import Node @@ -123,11 +123,21 @@ def test_should_choice_convert_enum(): Table('translatedmodel', Base.metadata, column) graphene_type = convert_sqlalchemy_column(column) - assert issubclass(graphene_type, graphene.Enum) + assert isinstance(graphene_type, graphene.Enum) assert graphene_type._meta.name == 'TRANSLATEDMODEL_LANGUAGE' assert graphene_type._meta.description == 'Language' assert graphene_type._meta.enum.__members__['es'].value == 'Spanish' assert graphene_type._meta.enum.__members__['en'].value == 'English' + assert callable(graphene_type.serialize) + + class MockEnumType: + _value_lookup = { + 'Spanish': graphene_type._meta.enum.__members__['es'], + 'English': graphene_type._meta.enum.__members__['en'] + } + serialize = getattr(type(graphene_type), 'serialize') + assert serialize(MockEnumType(), Choice('es', 'Spanish')) == 'es' + assert serialize(MockEnumType(), Choice('fr', 'French')) is None def test_should_scalar_list_convert_list():