Serialization function should be applied on ChoiceType output

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)

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():