Skip to content

Commit afed25a

Browse files
committed
Improved enums comparison and getters. Fixed #444
1 parent aaf9e92 commit afed25a

File tree

3 files changed

+90
-1
lines changed

3 files changed

+90
-1
lines changed

docs/types/enums.rst

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,33 @@ Notes
5959
-----
6060

6161
``graphene.Enum`` uses |enum.Enum|_ internally (or a backport if
62-
that's not available) and can be used in the exact same way.
62+
that's not available) and can be used in a similar way, with the exception of
63+
member getters.
64+
65+
In the Python ``Enum`` implementation you can access a member by initing the Enum.
66+
67+
.. code:: python
68+
69+
from enum import Enum
70+
class Color(Enum):
71+
RED = 1
72+
GREEN = 2
73+
BLUE = 3
74+
75+
assert Color(1) == Color.RED
76+
77+
78+
However, in Graphene ``Enum`` you need to call get to have the same effect:
79+
80+
.. code:: python
81+
82+
from graphene import Enum
83+
class Color(Enum):
84+
RED = 1
85+
GREEN = 2
86+
BLUE = 3
87+
88+
assert Color.get(1) == Color.RED
6389
6490
.. |enum.Enum| replace:: ``enum.Enum``
6591
.. _enum.Enum: https://docs.python.org/3/library/enum.html

graphene/types/enum.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,12 @@
1313
from ..pyutils.enum import Enum as PyEnum
1414

1515

16+
def eq_enum(self, other):
17+
if isinstance(other, self.__class__):
18+
return self is other
19+
return self.value is other
20+
21+
1622
class EnumTypeMeta(type):
1723

1824
def __new__(cls, name, bases, attrs):
@@ -28,6 +34,7 @@ def __new__(cls, name, bases, attrs):
2834
enum=None,
2935
)
3036
if not options.enum:
37+
attrs['__eq__'] = eq_enum
3138
options.enum = PyEnum(cls.__name__, attrs)
3239

3340
new_attrs = OrderedDict(attrs, _meta=options, **options.enum.__members__)
@@ -36,11 +43,18 @@ def __new__(cls, name, bases, attrs):
3643
def __prepare__(name, bases, **kwargs): # noqa: N805
3744
return OrderedDict()
3845

46+
def get(cls, value):
47+
return cls._meta.enum(value)
48+
49+
def __getitem__(cls, value):
50+
return cls._meta.enum[value]
51+
3952
def __call__(cls, *args, **kwargs): # noqa: N805
4053
if cls is Enum:
4154
description = kwargs.pop('description', None)
4255
return cls.from_enum(PyEnum(*args, **kwargs), description=description)
4356
return super(EnumTypeMeta, cls).__call__(*args, **kwargs)
57+
# return cls._meta.enum(*args, **kwargs)
4458

4559
def from_enum(cls, enum, description=None): # noqa: N805
4660
meta_class = type('Meta', (object,), {'enum': enum, 'description': description})

graphene/types/tests/test_enum.py

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,3 +111,52 @@ class RGB(Enum):
111111
unmounted_field = unmounted.Argument()
112112
assert isinstance(unmounted_field, Argument)
113113
assert unmounted_field.type == RGB
114+
115+
116+
def test_enum_can_be_compared():
117+
class RGB(Enum):
118+
RED = 1
119+
GREEN = 2
120+
BLUE = 3
121+
122+
assert RGB.RED == 1
123+
assert RGB.GREEN == 2
124+
assert RGB.BLUE == 3
125+
126+
127+
def test_enum_can_be_initialzied():
128+
class RGB(Enum):
129+
RED = 1
130+
GREEN = 2
131+
BLUE = 3
132+
133+
assert RGB.get(1) == RGB.RED
134+
assert RGB.get(2) == RGB.GREEN
135+
assert RGB.get(3) == RGB.BLUE
136+
137+
138+
def test_enum_can_retrieve_members():
139+
class RGB(Enum):
140+
RED = 1
141+
GREEN = 2
142+
BLUE = 3
143+
144+
assert RGB['RED'] == RGB.RED
145+
assert RGB['GREEN'] == RGB.GREEN
146+
assert RGB['BLUE'] == RGB.BLUE
147+
148+
149+
def test_enum_to_enum_comparison_should_differ():
150+
class RGB1(Enum):
151+
RED = 1
152+
GREEN = 2
153+
BLUE = 3
154+
155+
class RGB2(Enum):
156+
RED = 1
157+
GREEN = 2
158+
BLUE = 3
159+
160+
assert RGB1.RED != RGB2.RED
161+
assert RGB1.GREEN != RGB2.GREEN
162+
assert RGB1.BLUE != RGB2.BLUE

0 commit comments

Comments
 (0)