Add support for Python enums in sqlalchemy_utils.ChoiceType (#240)

Author: metheoryt

graphene_sqlalchemy/converter.py

@@ -1,3 +1,5 @@
+from enum import EnumMeta
+
 from singledispatch import singledispatch
 from sqlalchemy import types
 from sqlalchemy.dialects import postgresql
@@ -163,7 +165,12 @@ def convert_enum_to_enum(type, column, registry=None):
 @convert_sqlalchemy_type.register(ChoiceType)
 def convert_choice_to_enum(type, column, registry=None):
     name = "{}_{}".format(column.table.name, column.name).upper()
-    return Enum(name, type.choices)
+    if isinstance(type.choices, EnumMeta):
+        # type.choices may be Enum/IntEnum, in ChoiceType both presented as EnumMeta
+        # do not use from_enum here because we can have more than one enum column in table
+        return Enum(name, list((v.name, v.value) for v in type.choices))
+    else:
+        return Enum(name, type.choices)
 
 
 @convert_sqlalchemy_type.register(ScalarListType)

graphene_sqlalchemy/tests/test_converter.py

@@ -145,6 +145,32 @@ def test_should_choice_convert_enum():
     assert graphene_type._meta.enum.__members__["en"].value == "English"
 
 
+def test_should_enum_choice_convert_enum():
+    class TestEnum(enum.Enum):
+        es = u"Spanish"
+        en = u"English"
+
+    field = get_field(ChoiceType(TestEnum, impl=types.String()))
+    graphene_type = field.type
+    assert issubclass(graphene_type, graphene.Enum)
+    assert graphene_type._meta.name == "MODEL_COLUMN"
+    assert graphene_type._meta.enum.__members__["es"].value == "Spanish"
+    assert graphene_type._meta.enum.__members__["en"].value == "English"
+
+
+def test_should_intenum_choice_convert_enum():
+    class TestEnum(enum.IntEnum):
+        one = 1
+        two = 2
+
+    field = get_field(ChoiceType(TestEnum, impl=types.String()))
+    graphene_type = field.type
+    assert issubclass(graphene_type, graphene.Enum)
+    assert graphene_type._meta.name == "MODEL_COLUMN"
+    assert graphene_type._meta.enum.__members__["one"].value == 1
+    assert graphene_type._meta.enum.__members__["two"].value == 2
+
+
 def test_should_columproperty_convert():
     field = get_field_from_column(column_property(
         select([func.sum(func.cast(id, types.Integer))]).where(id == 1)