Merge pull request #446 from graphql-python/features/improved-enums

Improved enums comparison and getters

Author: syrusakbary

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

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})

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