Improved plugin execution

Author: syrusakbary

graphene/core/schema.py

@@ -119,24 +119,9 @@ def types(self):
         return self._types_names
 
     def execute(self, request='', root=None, args=None, **kwargs):
-        executor = kwargs
-        executor['root'] = root
-        executor['args'] = args
-        executor['schema'] = self.schema
-        executor['request'] = request
-        contexts = []
-        for plugin in self.plugins:
-            if not hasattr(plugin, 'context_execution'):
-                continue
-            context = plugin.context_execution(executor)
-            executor = context.__enter__()
-            contexts.append((context, executor))
-        result = self.executor.execute(
-            **executor
-        )
-        for context, value in contexts[::-1]:
-            context.__exit__(None, None, None)
-        return result
+        kwargs = dict(kwargs, request=request, root=root, args=args, schema=self.schema)
+        with self.plugins.context_execution(**kwargs) as execute_kwargs:
+            return self.executor.execute(**execute_kwargs)
 
     def introspect(self):
         return self.execute(introspection_query).data

graphene/plugins/base.py

@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 from functools import partial, reduce
 
 
@@ -38,3 +39,15 @@ def __getattr__(self, name):
 
     def __contains__(self, name):
         return name in self.PLUGIN_FUNCTIONS
+
+    @contextmanager
+    def context_execution(self, **executor):
+        contexts = []
+        functions = self.get_plugin_functions('context_execution')
+        for f in functions:
+            context = f(executor)
+            executor = context.__enter__()
+            contexts.append((context, executor))
+        yield executor
+        for context, value in contexts[::-1]:
+            context.__exit__(None, None, None)