diff --git a/graphene/new_types/tests/test_definition.py b/graphene/new_types/tests/test_definition.py index 7e2811c7..13a0befc 100644 --- a/graphene/new_types/tests/test_definition.py +++ b/graphene/new_types/tests/test_definition.py @@ -182,6 +182,28 @@ def test_includes_types_in_union(): assert schema.get_type_map()['SomeType'].graphene_type is SomeType +def test_maps_enum(): + class SomeType(ObjectType): + a = String() + + class OtherType(ObjectType): + b = String() + + class MyUnion(Union): + class Meta: + types = (SomeType, OtherType) + + class Query(ObjectType): + union = Field(MyUnion) + + schema = Schema( + query=Query, + ) + + assert schema.get_type_map()['OtherType'].graphene_type is OtherType + assert schema.get_type_map()['SomeType'].graphene_type is SomeType + + def test_includes_interfaces_subtypes_in_the_type_map(): class SomeInterface(Interface): f = Int() diff --git a/graphene/new_types/tests/test_typemap.py b/graphene/new_types/tests/test_typemap.py new file mode 100644 index 00000000..7a47edbf --- /dev/null +++ b/graphene/new_types/tests/test_typemap.py @@ -0,0 +1,36 @@ +import pytest + +from ..objecttype import ObjectType +from ..union import Union +from ..enum import Enum +from ..typemap import TypeMap + +from graphql.type import GraphQLEnumType, GraphQLEnumValue + + +def test_enum(): + class MyEnum(Enum): + '''Description''' + foo = 1 + bar = 2 + + @property + def description(self): + return 'Description {}={}'.format(self.name, self.value) + + @property + def deprecation_reason(self): + if self == MyEnum.foo: + return 'Is deprecated' + + typemap = TypeMap([MyEnum]) + assert 'MyEnum' in typemap + graphql_enum = typemap['MyEnum'] + assert isinstance(graphql_enum, GraphQLEnumType) + assert graphql_enum.name == 'MyEnum' + assert graphql_enum.description == 'Description' + values = graphql_enum.get_values() + assert values == [ + GraphQLEnumValue(name='foo', value=1, description='Description foo=1', deprecation_reason='Is deprecated'), + GraphQLEnumValue(name='bar', value=2, description='Description bar=2'), + ] diff --git a/graphene/new_types/typemap.py b/graphene/new_types/typemap.py index a352904e..e981dcf4 100644 --- a/graphene/new_types/typemap.py +++ b/graphene/new_types/typemap.py @@ -8,15 +8,17 @@ from .interface import Interface from .union import Union from .inputobjecttype import InputObjectType from .structures import List, NonNull +from .enum import Enum from .scalars import Scalar, String, Boolean, Int, Float, ID from graphql import GraphQLString, GraphQLField, GraphQLList, GraphQLBoolean, GraphQLInt, GraphQLFloat, GraphQLID, GraphQLNonNull, GraphQLInputObjectField +from graphql.type import GraphQLEnumValue def is_graphene_type(_type): if isinstance(_type, (List, NonNull)): return True - if inspect.isclass(_type) and issubclass(_type, (ObjectType, InputObjectType, Scalar, Interface, Union)): + if inspect.isclass(_type) and issubclass(_type, (ObjectType, InputObjectType, Scalar, Interface, Union, Enum)): return True @@ -48,6 +50,8 @@ class TypeMap(GraphQLTypeMap): return cls.construct_interface(map, type) if issubclass(type, Scalar): return cls.construct_scalar(map, type) + if issubclass(type, Enum): + return cls.construct_enum(map, type) if issubclass(type, Union): return cls.construct_union(map, type) return map @@ -76,6 +80,25 @@ class TypeMap(GraphQLTypeMap): ) return map + @classmethod + def construct_enum(cls, map, type): + from ..generators.definitions import GrapheneEnumType + values = OrderedDict() + for name, value in type._meta.enum.__members__.items(): + values[name] = GraphQLEnumValue( + name=name, + value=value.value, + description=getattr(value, 'description', None), + deprecation_reason=getattr(value, 'deprecation_reason', None) + ) + map[type._meta.name] = GrapheneEnumType( + graphene_type=type, + values=values, + name=type._meta.name, + description=type._meta.description, + ) + return map + @classmethod def construct_objecttype(cls, map, type): from ..generators.definitions import GrapheneObjectType