Merge pull request #586 from ekampf/feature/enum_value_descriptions

Support from_enum description lambda

Author: syrusakbary

graphene/types/enum.py

@@ -21,6 +21,7 @@ def eq_enum(self, other):
 
 class EnumOptions(BaseOptions):
     enum = None  # type: Enum
+    deprecation_reason = None
 
 
 class EnumMeta(SubclassWithMeta_Meta):
@@ -45,8 +46,14 @@ def __call__(cls, *args, **kwargs):  # noqa: N805
         return super(EnumMeta, cls).__call__(*args, **kwargs)
         # return cls._meta.enum(*args, **kwargs)
 
-    def from_enum(cls, enum, description=None):  # noqa: N805
-        meta_class = type('Meta', (object,), {'enum': enum, 'description': description})
+    def from_enum(cls, enum, description=None, deprecation_reason=None):  # noqa: N805
+        description = description or enum.__doc__
+        meta_dict = {
+            'enum': enum,
+            'description': description,
+            'deprecation_reason': deprecation_reason
+        }
+        meta_class = type('Meta', (object,), meta_dict)
         return type(meta_class.enum.__name__, (Enum,), {'Meta': meta_class})
 
 
@@ -56,8 +63,10 @@ class Enum(six.with_metaclass(EnumMeta, UnmountedType, BaseType)):
     def __init_subclass_with_meta__(cls, enum=None, **options):
         _meta = EnumOptions(cls)
         _meta.enum = enum or cls.__enum__
+        _meta.deprecation_reason = options.pop('deprecation_reason', None)
         for key, value in _meta.enum.__members__.items():
             setattr(cls, key, value)
+
         super(Enum, cls).__init_subclass_with_meta__(_meta=_meta, **options)
 
     @classmethod

graphene/types/tests/test_enum.py

@@ -1,3 +1,6 @@
+import six
+
+from ..schema import Schema, ObjectType
 from ..argument import Argument
 from ..enum import Enum, PyEnum
 from ..field import Field
@@ -66,6 +69,57 @@ def test_enum_from_builtin_enum():
     assert RGB.BLUE
 
 
+def test_enum_from_builtin_enum_accepts_lambda_description():
+    def custom_description(value):
+        if not value:
+            return "StarWars Episodes"
+
+        return 'New Hope Episode' if value == Episode.NEWHOPE else 'Other'
+
+    def custom_deprecation_reason(value):
+        return 'meh' if value == Episode.NEWHOPE else None
+
+    PyEpisode = PyEnum('PyEpisode', 'NEWHOPE,EMPIRE,JEDI')
+    Episode = Enum.from_enum(PyEpisode, description=custom_description, deprecation_reason=custom_deprecation_reason)
+
+    class Query(ObjectType):
+        foo = Episode()
+
+    schema = Schema(query=Query)
+
+    GraphQLPyEpisode = schema._type_map['PyEpisode'].values
+
+    assert schema._type_map['PyEpisode'].description == "StarWars Episodes"
+    assert GraphQLPyEpisode[0].name == 'NEWHOPE' and GraphQLPyEpisode[0].description == 'New Hope Episode'
+    assert GraphQLPyEpisode[1].name == 'EMPIRE' and GraphQLPyEpisode[1].description == 'Other'
+    assert GraphQLPyEpisode[2].name == 'JEDI' and GraphQLPyEpisode[2].description == 'Other'
+
+    assert GraphQLPyEpisode[0].name == 'NEWHOPE' and GraphQLPyEpisode[0].deprecation_reason == 'meh'
+    assert GraphQLPyEpisode[1].name == 'EMPIRE' and GraphQLPyEpisode[1].deprecation_reason == None
+    assert GraphQLPyEpisode[2].name == 'JEDI' and GraphQLPyEpisode[2].deprecation_reason == None
+
+
+def test_enum_from_python3_enum_uses_enum_doc():
+    if not six.PY3:
+        return
+
+    from enum import Enum as PyEnum
+
+    class Color(PyEnum):
+        """This is the description"""
+        RED = 1
+        GREEN = 2
+        BLUE = 3
+
+    RGB = Enum.from_enum(Color)
+    assert RGB._meta.enum == Color
+    assert RGB._meta.description == "This is the description"
+    assert RGB
+    assert RGB.RED
+    assert RGB.GREEN
+    assert RGB.BLUE
+
+
 def test_enum_value_from_class():
     class RGB(Enum):
         RED = 1

graphene/types/typemap.py

@@ -127,16 +127,27 @@ def construct_scalar(self, map, type):
     def construct_enum(self, map, type):
         values = OrderedDict()
         for name, value in type._meta.enum.__members__.items():
+            description = getattr(value, 'description', None)
+            deprecation_reason = getattr(value, 'deprecation_reason', None)
+            if not description and callable(type._meta.description):
+                description = type._meta.description(value)
+
+            if not deprecation_reason and callable(type._meta.deprecation_reason):
+                deprecation_reason = type._meta.deprecation_reason(value)
+
             values[name] = GraphQLEnumValue(
                 name=name,
                 value=value.value,
-                description=getattr(value, 'description', None),
-                deprecation_reason=getattr(value, 'deprecation_reason', None))
+                description=description,
+                deprecation_reason=deprecation_reason)
+
+        type_description = type._meta.description(None) if callable(type._meta.description) else type._meta.description
+
         return GrapheneEnumType(
             graphene_type=type,
             values=values,
             name=type._meta.name,
-            description=type._meta.description, )
+            description=type_description, )
 
     def construct_objecttype(self, map, type):
         if type._meta.name in map: