diff --git a/graphene/types/enum.py b/graphene/types/enum.py index a328cc93..4e3f2ac2 100644 --- a/graphene/types/enum.py +++ b/graphene/types/enum.py @@ -21,6 +21,7 @@ EnumType = type(PyEnum) class EnumOptions(BaseOptions): enum = None # type: Enum + deprecation_reason = None class EnumMeta(SubclassWithMeta_Meta): @@ -45,8 +46,14 @@ class EnumMeta(SubclassWithMeta_Meta): 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 diff --git a/graphene/types/tests/test_enum.py b/graphene/types/tests/test_enum.py index 42f00605..c4cf3b85 100644 --- a/graphene/types/tests/test_enum.py +++ b/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 diff --git a/graphene/types/typemap.py b/graphene/types/typemap.py index 3b389cd6..d4a4d157 100644 --- a/graphene/types/typemap.py +++ b/graphene/types/typemap.py @@ -127,16 +127,27 @@ class TypeMap(GraphQLTypeMap): 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: