Improved Enum implementation

graphene/pyutils/init_subclass.py

@@ -1,15 +1,19 @@
+import six
 
-class InitSubclassMeta(type):
-    """Metaclass that implements PEP 487 protocol"""
-    def __new__(cls, name, bases, ns):
-        __init_subclass__ = ns.pop('__init_subclass__', None)
-        if __init_subclass__:
-            __init_subclass__ = classmethod(__init_subclass__)
-            ns['__init_subclass__'] = __init_subclass__
-        return type.__new__(cls, name, bases, ns)
+if six.PY2:
+    class InitSubclassMeta(type):
+        """Metaclass that implements PEP 487 protocol"""
+        def __new__(cls, name, bases, ns):
+            __init_subclass__ = ns.pop('__init_subclass__', None)
+            if __init_subclass__:
+                __init_subclass__ = classmethod(__init_subclass__)
+                ns['__init_subclass__'] = __init_subclass__
+            return super(InitSubclassMeta, cls).__new__(cls, name, bases, ns)
 
-    def __init__(cls, name, bases, ns):
-        super(InitSubclassMeta, cls).__init__(name, bases, ns)
-        super_class = super(cls, cls)
-        if hasattr(super_class, '__init_subclass__'):
-            super_class.__init_subclass__.__func__(cls)
+        def __init__(cls, name, bases, ns):
+            super(InitSubclassMeta, cls).__init__(name, bases, ns)
+            super_class = super(cls, cls)
+            if hasattr(super_class, '__init_subclass__'):
+                super_class.__init_subclass__.__func__(cls)
+else:
+    InitSubclassMeta = type

graphene/types/enum.py

@@ -2,10 +2,9 @@ from collections import OrderedDict
 
 import six
 
-from ..utils.is_base_type import is_base_type
-from ..utils.trim_docstring import trim_docstring
-from .options import Options
+from .base import BaseOptions, BaseType
 from .unmountedtype import UnmountedType
+from graphene.pyutils.init_subclass import InitSubclassMeta
 
 try:
     from enum import Enum as PyEnum
@@ -19,59 +18,43 @@ def eq_enum(self, other):
     return self.value is other
 
 
-class EnumTypeMeta(type):
+EnumType = type(PyEnum)
 
-    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=trim_docstring(attrs.get('__doc__')),
-            enum=None,
-        )
-        if not options.enum:
-            attrs['__eq__'] = eq_enum
-            options.enum = PyEnum(cls.__name__, attrs)
+class EnumOptions(BaseOptions):
+    enum = None  # type: Enum
 
-        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()
 
+class EnumMeta(InitSubclassMeta):
     def get(cls, value):
         return cls._meta.enum(value)
 
     def __getitem__(cls, value):
         return cls._meta.enum[value]
 
+    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)
+        return super(EnumMeta, 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})
         return type(meta_class.enum.__name__, (Enum,), {'Meta': meta_class})
 
-    def __str__(cls):  # noqa: N805
-        return cls._meta.name
 
-
-class Enum(six.with_metaclass(EnumTypeMeta, UnmountedType)):
-    '''
-    Enum Type Definition
-
-    Some leaf values of requests and input values are Enums. GraphQL serializes
-    Enum values as strings, however internally Enums can be represented by any
-    kind of type, often integers.
-    '''
+class Enum(six.with_metaclass(EnumMeta, UnmountedType, BaseType)):
+    @classmethod
+    def __init_subclass_with_meta__(cls, enum=None, **options):
+        _meta = EnumOptions(cls)
+        _meta.enum = enum or PyEnum(cls.__name__, dict(cls.__dict__, __eq__=eq_enum))
+        for key, value in _meta.enum.__members__.items():
+            setattr(cls, key, value)
+        super(Enum, cls).__init_subclass_with_meta__(_meta=_meta, **options)
 
     @classmethod
     def get_type(cls):