Merge pull request #100 from graphql-python/features/enum

Add Enum type

.coveragerc

@@ -0,0 +1,2 @@
+[run]
+omit = graphene/utils/enum.py,graphene/contrib/django/debug/sql/tracking.py,*/tests/*

.travis.yml

@@ -23,7 +23,7 @@ before_install:
 install:
 - |
   if [ "$TEST_TYPE" = build ]; then
-    pip install --download-cache $HOME/.cache/pip/ pytest pytest-cov coveralls six pytest-django django-filter
+    pip install --download-cache $HOME/.cache/pip/ pytest pytest-cov coveralls six pytest-django django-filter sqlalchemy_utils
     pip install --download-cache $HOME/.cache/pip/ -e .[django]
     pip install --download-cache $HOME/.cache/pip/ -e .[sqlalchemy]
     pip install django==$DJANGO_VERSION

docs/config.toml

@@ -14,6 +14,7 @@ ga = "UA-12613282-7"
 	  "/docs/objecttypes/",
 	  "/docs/mutations/",
 	  "/docs/basic-types/",
+	  "/docs/enums/",
 	  "/docs/relay/",
 	]
 

docs/pages/docs/enums.md

@@ -0,0 +1,33 @@
+---
+title: Enums
+description: Walkthrough Enums
+---
+
+# Enums
+
+A `Enum` is a special `GraphQL` type that represents a set of symbolic names (members) bound to unique, constant values.
+
+## Enum definition
+
+You can create an `Enum` using classes:
+
+```python
+import graphene
+
+class Episode(graphene.Enum):
+    NEWHOPE = 4
+    EMPIRE = 5
+    JEDI = 6
+```
+
+But also using instances of Enum:
+
+```python
+Episode = graphene.Enum('Episode', [('NEWHOPE', 4), ('EMPIRE', 5), ('JEDI', 6)])
+```
+
+## Notes
+
+Internally, `graphene.Enum` uses [`enum.Enum`](https://docs.python.org/3/library/enum.html) Python implementation if available, or a backport if not.
+
+So you can use it in the same way as you would do with Python `enum.Enum`.

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/contrib/django/converter.py

@@ -1,12 +1,22 @@
 from django.db import models
 
 from ...core.types.scalars import ID, Boolean, Float, Int, String
+from ...core.classtypes.enum import Enum
 from .compat import RelatedObject, UUIDField
 from .utils import get_related_model, import_single_dispatch
 
 singledispatch = import_single_dispatch()
 
 
+def convert_django_field_with_choices(field):
+    choices = getattr(field, 'choices', None)
+    if choices:
+        meta = field.model._meta
+        name = '{}_{}_{}'.format(meta.app_label, meta.object_name, field.name)
+        return Enum(name.upper(), choices, description=field.help_text)
+    return convert_django_field(field)
+
+
 @singledispatch
 def convert_django_field(field):
     raise Exception(

graphene/contrib/django/tests/test_converter.py

@@ -2,7 +2,8 @@
 from py.test import raises
 
 import graphene
-from graphene.contrib.django.converter import convert_django_field
+from graphene.contrib.django.converter import (
+    convert_django_field, convert_django_field_with_choices)
 from graphene.contrib.django.fields import (ConnectionOrListField,
                                             DjangoModelField)
 
@@ -86,6 +87,26 @@ def test_should_nullboolean_convert_boolean():
     assert field.required is False
 
 
+def test_field_with_choices_convert_enum():
+    field = models.CharField(help_text='Language', choices=(
+        ('es', 'Spanish'),
+        ('en', 'English')
+    ))
+
+    class TranslatedModel(models.Model):
+        language = field
+
+        class Meta:
+            app_label = 'test'
+
+    graphene_type = convert_django_field_with_choices(field)
+    assert issubclass(graphene_type, graphene.Enum)
+    assert graphene_type._meta.type_name == 'TEST_TRANSLATEDMODEL_LANGUAGE'
+    assert graphene_type._meta.description == 'Language'
+    assert graphene_type.__enum__.__members__['es'].value == 'Spanish'
+    assert graphene_type.__enum__.__members__['en'].value == 'English'
+
+
 def test_should_float_convert_float():
     assert_conversion(models.FloatField, graphene.Float)
 

graphene/contrib/django/types.py

@@ -5,7 +5,7 @@
 
 from ...core.classtypes.objecttype import ObjectType, ObjectTypeMeta
 from ...relay.types import Connection, Node, NodeMeta
-from .converter import convert_django_field
+from .converter import convert_django_field_with_choices
 from .options import DjangoOptions
 from .utils import get_reverse_fields
 
@@ -29,7 +29,7 @@ def construct_fields(cls):
                 # We skip this field if we specify only_fields and is not
                 # in there. Or when we exclude this field in exclude_fields
                 continue
-            converted_field = convert_django_field(field)
+            converted_field = convert_django_field_with_choices(field)
             cls.add_to_class(field.name, converted_field)
 
     def construct(cls, *args, **kwargs):

graphene/contrib/sqlalchemy/converter.py

@@ -2,7 +2,13 @@
 
 from sqlalchemy import types
 from sqlalchemy.orm import interfaces
+try:
+    from sqlalchemy_utils.types.choice import ChoiceType
+except ImportError:
+    class ChoiceType(object):
+        pass
 
+from ...core.classtypes.enum import Enum
 from ...core.types.scalars import ID, Boolean, Float, Int, String
 from .fields import ConnectionOrListField, SQLAlchemyModelField
 
@@ -59,3 +65,9 @@ def convert_column_to_boolean(type, column):
 @convert_sqlalchemy_type.register(types.Numeric)
 def convert_column_to_float(type, column):
     return Float(description=column.doc)
+
+
+@convert_sqlalchemy_type.register(ChoiceType)
+def convert_column_to_enum(type, column):
+    name = '{}_{}'.format(column.table.name, column.name).upper()
+    return Enum(name, type.choices, description=column.doc)

graphene/contrib/sqlalchemy/tests/test_converter.py

@@ -5,7 +5,9 @@
                                                    convert_sqlalchemy_relationship)
 from graphene.contrib.sqlalchemy.fields import (ConnectionOrListField,
                                                 SQLAlchemyModelField)
-from sqlalchemy import Column, types
+from sqlalchemy import Table, Column, types
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy_utils.types.choice import ChoiceType
 
 from .models import Article, Pet, Reporter
 
@@ -85,6 +87,23 @@ def test_should_numeric_convert_float():
     assert_column_conversion(types.Numeric(), graphene.Float)
 
 
+def test_should_choice_convert_enum():
+    TYPES = [
+        (u'es', u'Spanish'),
+        (u'en', u'English')
+    ]
+    column = Column(ChoiceType(TYPES), doc='Language', name='language')
+    Base = declarative_base()
+
+    Table('translatedmodel', Base.metadata, column)
+    graphene_type = convert_sqlalchemy_column(column)
+    assert issubclass(graphene_type, graphene.Enum)
+    assert graphene_type._meta.type_name == 'TRANSLATEDMODEL_LANGUAGE'
+    assert graphene_type._meta.description == 'Language'
+    assert graphene_type.__enum__.__members__['es'].value == 'Spanish'
+    assert graphene_type.__enum__.__members__['en'].value == 'English'
+
+
 def test_should_manytomany_convert_connectionorlist():
     graphene_type = convert_sqlalchemy_relationship(Reporter.pets.property)
     assert isinstance(graphene_type, ConnectionOrListField)

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