Merge pull request #173 from arunsureshkumar/support-monogengine-field-enum

BugFix: import condition for Enum Field updated based on Python version

Author: abawchen

graphene_mongo/converter.py

@@ -429,9 +429,11 @@ def dynamic_type():
     return graphene.Dynamic(dynamic_type)
 
 
-if sys.version_info[0] > 3.5:
+if sys.version_info >= (3, 6):
     @convert_mongoengine_field.register(mongoengine.EnumField)
     def convert_field_to_enum(field, registry=None):
-        return graphene.Field(graphene.Enum.from_enum(field._enum_cls),
-                              description=get_field_description(field, registry), required=field.required
-                              )
+        if not registry.check_enum_already_exist(field._enum_cls):
+            registry.register_enum(field._enum_cls)
+        _type = registry.get_type_for_enum(field._enum_cls)
+        return graphene.Field(_type,
+                              description=get_field_description(field, registry), required=field.required)

graphene_mongo/registry.py

@@ -1,15 +1,19 @@
+from graphene import Enum
+
+
 class Registry(object):
     def __init__(self):
         self._registry = {}
         self._registry_string_map = {}
+        self._registry_enum = {}
 
     def register(self, cls):
         from .types import GrapheneMongoengineObjectTypes
 
         assert issubclass(
             cls,
             GrapheneMongoengineObjectTypes
-            ), 'Only Mongoengine object types can be registered, received "{}"'.format(
+        ), 'Only Mongoengine object types can be registered, received "{}"'.format(
             cls.__name__
         )
         assert cls._meta.registry == self, "Registry for a Model have to match."
@@ -20,9 +24,22 @@ def register(self, cls):
         for model, cls in self._registry.items():
             cls.rescan_fields()
 
+    def register_enum(self, cls):
+        from enum import EnumMeta
+        assert type(cls) == EnumMeta, 'Only EnumMeta can be registered, received "{}"'.format(
+            cls.__name__
+        )
+        self._registry_enum[cls] = Enum.from_enum(cls)
+
     def get_type_for_model(self, model):
         return self._registry.get(model)
 
+    def check_enum_already_exist(self, cls):
+        return cls in self._registry_enum
+
+    def get_type_for_enum(self, cls):
+        return self._registry_enum.get(cls)
+
 
 registry = None
 inputs_registry = None