Added Enum

graphene/new_types/enum.py

@@ -0,0 +1,51 @@
+from collections import OrderedDict
+
+import six
+
+from ..generators import generate_enum
+from ..utils.is_base_type import is_base_type
+from .options import Options
+from .unmountedtype import UnmountedType
+
+try:
+    from enum import Enum as PyEnum
+except ImportError:
+    from ..utils.enum import Enum as PyEnum
+
+
+class EnumTypeMeta(type):
+
+    def __new__(cls, name, bases, attrs):
+        # Also ensure initialization is only performed for subclasses of Model
+        # (excluding Model class itself).
+        if not is_base_type(bases, EnumTypeMeta):
+            return type.__new__(cls, name, bases, attrs)
+
+        options = Options(
+            attrs.pop('Meta', None),
+            name=name,
+            description=attrs.get('__doc__'),
+            enum=None,
+        )
+        if not options.enum:
+            options.enum = PyEnum(cls.__name__, attrs)
+
+        new_attrs = OrderedDict(attrs, _meta=options, **options.enum.__members__)
+        return type.__new__(cls, name, bases, new_attrs)
+
+    def __prepare__(name, bases, **kwargs):  # noqa: N805
+        return OrderedDict()
+
+    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)
+
+    def from_enum(cls, enum, description=None):
+        meta_class = type('Meta', (object,), {'enum': enum, 'description': description})
+        return type(meta_class.enum.__name__, (Enum,), {'Meta': meta_class})
+
+
+class Enum(six.with_metaclass(EnumTypeMeta, UnmountedType)):
+    pass

graphene/new_types/tests/test_enum.py

@@ -0,0 +1,73 @@
+from ..enum import Enum, PyEnum
+
+
+def test_enum_construction():
+    class RGB(Enum):
+        '''Description'''
+        RED = 1
+        GREEN = 2
+        BLUE = 3
+
+        @property
+        def description(self):
+            return "Description {}".format(self.name)
+
+    assert RGB._meta.name == 'RGB'
+    assert RGB._meta.description == 'Description'
+
+    values = RGB._meta.enum.__members__.values()
+    assert sorted([v.name for v in values]) == [
+        'BLUE',
+        'GREEN',
+        'RED'
+    ]
+    assert sorted([v.description for v in values]) == [
+        'Description BLUE',
+        'Description GREEN',
+        'Description RED'
+    ]
+
+
+def test_enum_construction_meta():
+    class RGB(Enum):
+        class Meta:
+            name = 'RGBEnum'
+            description = 'Description'
+        RED = 1
+        GREEN = 2
+        BLUE = 3
+
+    assert RGB._meta.name == 'RGBEnum'
+    assert RGB._meta.description == 'Description'
+
+
+def test_enum_instance_construction():
+    RGB = Enum('RGB', 'RED,GREEN,BLUE')
+
+    values = RGB._meta.enum.__members__.values()
+    assert sorted([v.name for v in values]) == [
+        'BLUE',
+        'GREEN',
+        'RED'
+    ]
+
+
+def test_enum_from_builtin_enum():
+    PyRGB = PyEnum('RGB', 'RED,GREEN,BLUE')
+
+    RGB = Enum.from_enum(PyRGB)
+    assert RGB._meta.enum == PyRGB
+    assert RGB.RED
+    assert RGB.GREEN
+    assert RGB.BLUE
+
+
+def test_enum_value_from_class():
+    class RGB(Enum):
+        RED = 1
+        GREEN = 2
+        BLUE = 3
+
+    assert RGB.RED.value == 1
+    assert RGB.GREEN.value == 2
+    assert RGB.BLUE.value == 3