Improved enums comparison and getters. Fixed #444

This commit is contained in:
Syrus Akbary 2017-04-07 22:02:16 -07:00
parent aaf9e92d24
commit afed25a180
3 changed files with 90 additions and 1 deletions

View File

@ -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

View File

@ -13,6 +13,12 @@ except ImportError:
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 @@ class EnumTypeMeta(type):
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 @@ class EnumTypeMeta(type):
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})

View File

@ -111,3 +111,52 @@ def test_enum_value_as_unmounted_argument():
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