From c30d8d02316fe3776a99cb32d4e0523f755f923d Mon Sep 17 00:00:00 2001
From: Jonathan Kim <jkimbo@gmail.com>
Date: Sat, 14 Mar 2020 16:46:19 +0000
Subject: [PATCH] Move enum middleware into compat module

---
 graphene/compat/__init__.py       |  0
 graphene/compat/middleware.py     | 28 ++++++++++++++++++++++++++++
 graphene/types/tests/test_enum.py | 29 ++---------------------------
 graphene/types/utils.py           |  7 +++++++
 4 files changed, 37 insertions(+), 27 deletions(-)
 create mode 100644 graphene/compat/__init__.py
 create mode 100644 graphene/compat/middleware.py

diff --git a/graphene/compat/__init__.py b/graphene/compat/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/graphene/compat/middleware.py b/graphene/compat/middleware.py
new file mode 100644
index 00000000..d5b1f682
--- /dev/null
+++ b/graphene/compat/middleware.py
@@ -0,0 +1,28 @@
+from graphql import OperationType
+
+from ..utils.str_converters import to_snake_case
+from ..types.definitions import GrapheneEnumType
+from ..types.utils import get_underlying_type
+
+
+def enum_value_convertor_middleware(next, root, info, **args):
+    """
+    Compatibility middleware for upgrading to v3:
+
+    Convert enums to their values for mutation inputs, like the behaviour in v2
+    """
+    operation = info.operation.operation
+    if operation == OperationType.MUTATION:
+        input_arguments = info.parent_type.fields[info.field_name].args
+        for arg_name, arg in input_arguments.items():
+            _type = get_underlying_type(arg.type)
+            if isinstance(_type, GrapheneEnumType):
+                # Convert inputs to value
+                arg_name = to_snake_case(arg_name)
+                input_value = args.get(arg_name, None)
+                if input_value and isinstance(
+                    input_value, _type.graphene_type._meta.enum
+                ):
+                    args[arg_name] = args[arg_name].value
+
+    return next(root, info, **args)
diff --git a/graphene/types/tests/test_enum.py b/graphene/types/tests/test_enum.py
index a57ed06e..95e712b1 100644
--- a/graphene/types/tests/test_enum.py
+++ b/graphene/types/tests/test_enum.py
@@ -1,9 +1,7 @@
 from textwrap import dedent
-from graphql import OperationType
 
-from graphene.utils.str_converters import to_snake_case
+from graphene.compat.middleware import enum_value_convertor_middleware
 from ..argument import Argument
-from ..definitions import GrapheneEnumType
 from ..enum import Enum, PyEnum
 from ..field import Field
 from ..inputfield import InputField
@@ -398,12 +396,6 @@ def test_enum_mutation():
     assert results.data["setFavColor"]["favColor"] == Color.RED.name
 
 
-def get_underlying_type(_type):
-    while hasattr(_type, "of_type"):
-        _type = _type.of_type
-    return _type
-
-
 def test_enum_mutation_compat():
     from enum import Enum as PyEnum
 
@@ -436,23 +428,6 @@ def test_enum_mutation_compat():
     class MyMutations(ObjectType):
         set_fav_color = SetFavColor.Field()
 
-    def enum_compat_middleware(next, root, info, **args):
-        operation = info.operation.operation
-        if operation == OperationType.MUTATION:
-            input_arguments = info.parent_type.fields[info.field_name].args
-            for arg_name, arg in input_arguments.items():
-                _type = get_underlying_type(arg.type)
-                if isinstance(_type, GrapheneEnumType):
-                    # Convert inputs to value
-                    arg_name = to_snake_case(arg_name)
-                    input_value = args.get(arg_name, None)
-                    if input_value and isinstance(
-                        input_value, _type.graphene_type._meta.enum
-                    ):
-                        args[arg_name] = args[arg_name].value
-
-        return next(root, info, **args)
-
     schema = Schema(query=Query, mutation=MyMutations)
 
     results = schema.execute(
@@ -461,7 +436,7 @@ def test_enum_mutation_compat():
                 favColor
             }
         }""",
-        middleware=[enum_compat_middleware],
+        middleware=[enum_value_convertor_middleware],
     )
     assert not results.errors
 
diff --git a/graphene/types/utils.py b/graphene/types/utils.py
index 3b195d69..1976448a 100644
--- a/graphene/types/utils.py
+++ b/graphene/types/utils.py
@@ -41,3 +41,10 @@ def get_type(_type):
     if inspect.isfunction(_type) or isinstance(_type, partial):
         return _type()
     return _type
+
+
+def get_underlying_type(_type):
+    """Get the underlying type even if it is wrapped in structures like NonNull"""
+    while hasattr(_type, "of_type"):
+        _type = _type.of_type
+    return _type