Added Enum class type

Author: syrusakbary

examples/starwars/schema.py

@@ -1,15 +1,13 @@
-from graphql.core.type import GraphQLEnumValue
-
 import graphene
 from graphene import resolve_only_args
 
 from .data import get_character, get_droid, get_hero, get_human
 
-Episode = graphene.Enum('Episode', dict(
-    NEWHOPE=GraphQLEnumValue(4),
-    EMPIRE=GraphQLEnumValue(5),
-    JEDI=GraphQLEnumValue(6)
-))
+
+class Episode(graphene.Enum):
+    NEWHOPE = 4
+    EMPIRE = 5
+    JEDI = 6
 
 
 class Character(graphene.Interface):

graphene/__init__.py

@@ -1,7 +1,3 @@
-from graphql.core.type import (
-    GraphQLEnumType as Enum
-)
-
 from graphene import signals
 
 from .core import (
@@ -11,6 +7,7 @@
     Interface,
     Mutation,
     Scalar,
+    Enum,
     InstanceType,
     LazyType,
     Argument,
@@ -58,6 +55,7 @@
     'Interface',
     'Mutation',
     'Scalar',
+    'Enum',
     'Field',
     'InputField',
     'StringField',

graphene/core/__init__.py

@@ -7,7 +7,8 @@
     InputObjectType,
     Interface,
     Mutation,
-    Scalar
+    Scalar,
+    Enum
 )
 
 from .types import (
@@ -42,5 +43,6 @@
     'Interface',
     'Mutation',
     'Scalar',
+    'Enum',
     'Field',
     'InputField']

graphene/core/classtypes/__init__.py

@@ -4,6 +4,7 @@
 from .objecttype import ObjectType
 from .options import Options
 from .scalar import Scalar
+from .enum import Enum
 from .uniontype import UnionType
 
 __all__ = [
@@ -13,4 +14,5 @@
     'ObjectType',
     'Options',
     'Scalar',
+    'Enum',
     'UnionType']

graphene/core/classtypes/enum.py

@@ -0,0 +1,45 @@
+import six
+from graphql.core.type import GraphQLEnumType, GraphQLEnumValue
+
+from .base import ClassTypeMeta, ClassType
+from ...utils.enum import Enum as PyEnum
+
+
+class EnumMeta(ClassTypeMeta):
+
+    def construct(cls, bases, attrs):
+        __enum__ = attrs.get('__enum__', None)
+        if not cls._meta.abstract and not __enum__:
+            __enum__ = PyEnum(cls._meta.type_name, attrs)
+            setattr(cls, '__enum__', __enum__)
+        if __enum__:
+            for k, v in __enum__.__members__.items():
+                attrs[k] = v.value
+        return super(EnumMeta, cls).construct(bases, attrs)
+
+    def __call__(cls, name, names=None, description=None):
+        attrs = {
+            '__enum__': PyEnum(name, names)
+        }
+        if description:
+            attrs['__doc__'] = description
+        return type(name, (Enum,), attrs)
+
+
+class Enum(six.with_metaclass(EnumMeta, ClassType)):
+
+    class Meta:
+        abstract = True
+
+    @classmethod
+    def internal_type(cls, schema):
+        if cls._meta.abstract:
+            raise Exception("Abstract Enum don't have a specific type.")
+
+        values = {k: GraphQLEnumValue(v.value) for k, v in cls.__enum__.__members__.items()}
+        # GraphQLEnumValue
+        return GraphQLEnumType(
+            cls._meta.type_name,
+            values=values,
+            description=cls._meta.description,
+        )

graphene/core/classtypes/tests/test_enum.py

@@ -0,0 +1,37 @@
+from graphql.core.type import GraphQLEnumType
+
+from graphene.core.schema import Schema
+
+from ..enum import Enum
+
+
+def test_enum():
+    class RGB(Enum):
+        '''RGB enum description'''
+        RED = 0
+        GREEN = 1
+        BLUE = 2
+
+    schema = Schema()
+
+    object_type = schema.T(RGB)
+    assert isinstance(object_type, GraphQLEnumType)
+    assert RGB._meta.type_name == 'RGB'
+    assert RGB._meta.description == 'RGB enum description'
+    assert RGB.RED == 0
+    assert RGB.GREEN == 1
+    assert RGB.BLUE == 2
+
+
+def test_enum_values():
+    RGB = Enum('RGB', dict(RED=0, GREEN=1, BLUE=2), description='RGB enum description')
+
+    schema = Schema()
+
+    object_type = schema.T(RGB)
+    assert isinstance(object_type, GraphQLEnumType)
+    assert RGB._meta.type_name == 'RGB'
+    assert RGB._meta.description == 'RGB enum description'
+    assert RGB.RED == 0
+    assert RGB.GREEN == 1
+    assert RGB.BLUE == 2