Merge pull request #139 from graphql-python/features/enum-improvements

Enum Improvements

graphene/contrib/django/converter.py

@@ -2,18 +2,24 @@ from django.db import models
 
 from ...core.types.scalars import ID, Boolean, Float, Int, String
 from ...core.classtypes.enum import Enum
+from ...utils import to_const
 from .compat import RelatedObject, UUIDField
 from .utils import get_related_model, import_single_dispatch
 
 singledispatch = import_single_dispatch()
 
 
+def convert_choices(choices):
+    for value, name in choices:
+        yield to_const(name), value
+
+
 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 Enum(name.upper(), list(convert_choices(choices)), description=field.help_text)
     return convert_django_field(field)
 
 

graphene/contrib/django/tests/models.py

@@ -30,6 +30,8 @@ class Article(models.Model):
         ('es', 'Spanish'),
         ('en', 'English')
     ], default='es')
+    importance = models.IntegerField('Importance', null=True, blank=True,
+                                     choices=[(1, u'Very important'), (2, u'Not as important')])
 
     def __str__(self):              # __unicode__ on Python 2
         return self.headline

graphene/contrib/django/tests/test_converter.py

@@ -103,8 +103,8 @@ def test_field_with_choices_convert_enum():
     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'
+    assert graphene_type.__enum__.__members__['SPANISH'].value == 'es'
+    assert graphene_type.__enum__.__members__['ENGLISH'].value == 'en'
 
 
 def test_should_float_convert_float():

graphene/core/classtypes/enum.py

@@ -2,6 +2,7 @@ import six
 from graphql.core.type import GraphQLEnumType, GraphQLEnumValue
 
 from .base import ClassTypeMeta, ClassType
+from ..types.base import MountedType
 from ...utils.enum import Enum as PyEnum
 
 
@@ -17,7 +18,12 @@ class EnumMeta(ClassTypeMeta):
                 attrs[k] = v.value
         return super(EnumMeta, cls).construct(bases, attrs)
 
-    def __call__(cls, name, names=None, description=None):
+    def __call__(cls, *args, **kwargs):
+        if cls is Enum:
+            return cls.create_enum(*args, **kwargs)
+        return super(EnumMeta, cls).__call__(*args, **kwargs)
+
+    def create_enum(cls, name, names=None, description=None):
         attrs = {
             '__enum__': PyEnum(name, names)
         }
@@ -26,7 +32,7 @@ class EnumMeta(ClassTypeMeta):
         return type(name, (Enum,), attrs)
 
 
-class Enum(six.with_metaclass(EnumMeta, ClassType)):
+class Enum(six.with_metaclass(EnumMeta, ClassType, MountedType)):
 
     class Meta:
         abstract = True

graphene/core/classtypes/tests/test_enum.py

@@ -3,6 +3,7 @@ from graphql.core.type import GraphQLEnumType
 from graphene.core.schema import Schema
 
 from ..enum import Enum
+from ..objecttype import ObjectType
 
 
 def test_enum():
@@ -35,3 +36,14 @@ def test_enum_values():
     assert RGB.RED == 0
     assert RGB.GREEN == 1
     assert RGB.BLUE == 2
+
+
+def test_enum_instance():
+    RGB = Enum('RGB', dict(RED=0, GREEN=1, BLUE=2))
+    RGB_field = RGB(description='RGB enum description')
+
+    class ObjectWithColor(ObjectType):
+        color = RGB_field
+
+    object_field = ObjectWithColor._meta.fields_map['color']
+    assert object_field.description == 'RGB enum description'

graphene/utils/__init__.py

@@ -1,4 +1,4 @@
-from .str_converters import to_camel_case, to_snake_case
+from .str_converters import to_camel_case, to_snake_case, to_const
 from .proxy_snake_dict import ProxySnakeDict
 from .caching import cached_property, memoize
 from .maybe_func import maybe_func
@@ -7,6 +7,6 @@ from .resolve_only_args import resolve_only_args
 from .lazylist import LazyList
 
 
-__all__ = ['to_camel_case', 'to_snake_case', 'ProxySnakeDict',
+__all__ = ['to_camel_case', 'to_snake_case', 'to_const', 'ProxySnakeDict',
            'cached_property', 'memoize', 'maybe_func', 'enum_to_graphql_enum',
            'resolve_only_args', 'LazyList']

graphene/utils/str_converters.py

@@ -15,3 +15,7 @@ def to_camel_case(snake_str):
 def to_snake_case(name):
     s1 = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', name)
     return re.sub('([a-z0-9])([A-Z])', r'\1_\2', s1).lower()
+
+
+def to_const(string):
+    return re.sub('[\W|^(?=\d)]+', '_', string).upper()

graphene/utils/tests/test_str_converter.py

@@ -1,4 +1,5 @@
-from ..str_converters import to_camel_case, to_snake_case
+# coding: utf-8
+from ..str_converters import to_camel_case, to_snake_case, to_const
 
 
 def test_snake_case():
@@ -15,3 +16,7 @@ def test_camel_case():
     assert to_camel_case('snakes_on_a_plane') == 'snakesOnAPlane'
     assert to_camel_case('snakes_on_a__plane') == 'snakesOnA_Plane'
     assert to_camel_case('i_phone_hysteria') == 'iPhoneHysteria'
+
+
+def test_to_const():
+    assert to_const('snakes $1. on a "#plane') == 'SNAKES_ON_A_PLANE'