Added enum support in typemap

graphene/new_types/tests/test_definition.py

@@ -182,6 +182,28 @@ def test_includes_types_in_union():
     assert schema.get_type_map()['SomeType'].graphene_type is SomeType
 
 
+def test_maps_enum():
+    class SomeType(ObjectType):
+        a = String()
+
+    class OtherType(ObjectType):
+        b = String()
+
+    class MyUnion(Union):
+        class Meta:
+            types = (SomeType, OtherType)
+
+    class Query(ObjectType):
+        union = Field(MyUnion)
+
+    schema = Schema(
+        query=Query,
+    )
+
+    assert schema.get_type_map()['OtherType'].graphene_type is OtherType
+    assert schema.get_type_map()['SomeType'].graphene_type is SomeType
+
+
 def test_includes_interfaces_subtypes_in_the_type_map():
     class SomeInterface(Interface):
         f = Int()

graphene/new_types/tests/test_typemap.py

@@ -0,0 +1,36 @@
+import pytest
+
+from ..objecttype import ObjectType
+from ..union import Union
+from ..enum import Enum
+from ..typemap import TypeMap
+
+from graphql.type import GraphQLEnumType, GraphQLEnumValue
+
+
+def test_enum():
+    class MyEnum(Enum):
+        '''Description'''
+        foo = 1
+        bar = 2
+
+        @property
+        def description(self):
+            return 'Description {}={}'.format(self.name, self.value)
+
+        @property
+        def deprecation_reason(self):
+            if self == MyEnum.foo:
+                return 'Is deprecated'
+
+    typemap = TypeMap([MyEnum])
+    assert 'MyEnum' in typemap
+    graphql_enum = typemap['MyEnum']
+    assert isinstance(graphql_enum, GraphQLEnumType)
+    assert graphql_enum.name == 'MyEnum'
+    assert graphql_enum.description == 'Description'
+    values = graphql_enum.get_values()
+    assert values == [
+        GraphQLEnumValue(name='foo', value=1, description='Description foo=1', deprecation_reason='Is deprecated'),
+        GraphQLEnumValue(name='bar', value=2, description='Description bar=2'),
+    ]

graphene/new_types/typemap.py

@@ -8,15 +8,17 @@ from .interface import Interface
 from .union import Union
 from .inputobjecttype import InputObjectType
 from .structures import List, NonNull
+from .enum import Enum
 from .scalars import Scalar, String, Boolean, Int, Float, ID
 
 from graphql import GraphQLString, GraphQLField, GraphQLList, GraphQLBoolean, GraphQLInt, GraphQLFloat, GraphQLID, GraphQLNonNull, GraphQLInputObjectField
+from graphql.type import GraphQLEnumValue
 
 
 def is_graphene_type(_type):
     if isinstance(_type, (List, NonNull)):
         return True
-    if inspect.isclass(_type) and issubclass(_type, (ObjectType, InputObjectType, Scalar, Interface, Union)):
+    if inspect.isclass(_type) and issubclass(_type, (ObjectType, InputObjectType, Scalar, Interface, Union, Enum)):
         return True
 
 
@@ -48,6 +50,8 @@ class TypeMap(GraphQLTypeMap):
             return cls.construct_interface(map, type)
         if issubclass(type, Scalar):
             return cls.construct_scalar(map, type)
+        if issubclass(type, Enum):
+            return cls.construct_enum(map, type)
         if issubclass(type, Union):
             return cls.construct_union(map, type)
         return map
@@ -76,6 +80,25 @@ class TypeMap(GraphQLTypeMap):
             )
         return map
 
+    @classmethod
+    def construct_enum(cls, map, type):
+        from ..generators.definitions import GrapheneEnumType
+        values = OrderedDict()
+        for name, value in type._meta.enum.__members__.items():
+            values[name] = GraphQLEnumValue(
+                name=name,
+                value=value.value,
+                description=getattr(value, 'description', None),
+                deprecation_reason=getattr(value, 'deprecation_reason', None)
+            )
+        map[type._meta.name] = GrapheneEnumType(
+            graphene_type=type,
+            values=values,
+            name=type._meta.name,
+            description=type._meta.description,
+        )
+        return map
+
     @classmethod
     def construct_objecttype(cls, map, type):
         from ..generators.definitions import GrapheneObjectType