diff --git a/docs/types/enums.rst b/docs/types/enums.rst index f0fefd11f..c2a184e45 100644 --- a/docs/types/enums.rst +++ b/docs/types/enums.rst @@ -59,7 +59,33 @@ Notes ----- ``graphene.Enum`` uses |enum.Enum|_ internally (or a backport if -that's not available) and can be used in the exact same way. +that's not available) and can be used in a similar way, with the exception of +member getters. + +In the Python ``Enum`` implementation you can access a member by initing the Enum. + +.. code:: python + + from enum import Enum + class Color(Enum): + RED = 1 + GREEN = 2 + BLUE = 3 + + assert Color(1) == Color.RED + + +However, in Graphene ``Enum`` you need to call get to have the same effect: + +.. code:: python + + from graphene import Enum + class Color(Enum): + RED = 1 + GREEN = 2 + BLUE = 3 + + assert Color.get(1) == Color.RED .. |enum.Enum| replace:: ``enum.Enum`` .. _enum.Enum: https://docs.python.org/3/library/enum.html diff --git a/graphene/types/enum.py b/graphene/types/enum.py index a6575cea5..029e6991a 100644 --- a/graphene/types/enum.py +++ b/graphene/types/enum.py @@ -13,6 +13,12 @@ from ..pyutils.enum import Enum as PyEnum +def eq_enum(self, other): + if isinstance(other, self.__class__): + return self is other + return self.value is other + + class EnumTypeMeta(type): def __new__(cls, name, bases, attrs): @@ -28,6 +34,7 @@ def __new__(cls, name, bases, attrs): enum=None, ) if not options.enum: + attrs['__eq__'] = eq_enum options.enum = PyEnum(cls.__name__, attrs) new_attrs = OrderedDict(attrs, _meta=options, **options.enum.__members__) @@ -36,11 +43,18 @@ def __new__(cls, name, bases, attrs): def __prepare__(name, bases, **kwargs): # noqa: N805 return OrderedDict() + def get(cls, value): + return cls._meta.enum(value) + + def __getitem__(cls, value): + return cls._meta.enum[value] + def __call__(cls, *args, **kwargs): # noqa: N805 if cls is Enum: description = kwargs.pop('description', None) return cls.from_enum(PyEnum(*args, **kwargs), description=description) return super(EnumTypeMeta, 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}) diff --git a/graphene/types/tests/test_enum.py b/graphene/types/tests/test_enum.py index a5a2d4c1f..6cd22bd98 100644 --- a/graphene/types/tests/test_enum.py +++ b/graphene/types/tests/test_enum.py @@ -111,3 +111,52 @@ class RGB(Enum): unmounted_field = unmounted.Argument() assert isinstance(unmounted_field, Argument) assert unmounted_field.type == RGB + + +def test_enum_can_be_compared(): + class RGB(Enum): + RED = 1 + GREEN = 2 + BLUE = 3 + + assert RGB.RED == 1 + assert RGB.GREEN == 2 + assert RGB.BLUE == 3 + + +def test_enum_can_be_initialzied(): + class RGB(Enum): + RED = 1 + GREEN = 2 + BLUE = 3 + + assert RGB.get(1) == RGB.RED + assert RGB.get(2) == RGB.GREEN + assert RGB.get(3) == RGB.BLUE + + +def test_enum_can_retrieve_members(): + class RGB(Enum): + RED = 1 + GREEN = 2 + BLUE = 3 + + assert RGB['RED'] == RGB.RED + assert RGB['GREEN'] == RGB.GREEN + assert RGB['BLUE'] == RGB.BLUE + + +def test_enum_to_enum_comparison_should_differ(): + class RGB1(Enum): + RED = 1 + GREEN = 2 + BLUE = 3 + + class RGB2(Enum): + RED = 1 + GREEN = 2 + BLUE = 3 + + assert RGB1.RED != RGB2.RED + assert RGB1.GREEN != RGB2.GREEN + assert RGB1.BLUE != RGB2.BLUE