Added get_operation_type method to GraphQLDocument

Author: syrusakbary

graphql/backend/base.py

@@ -1,4 +1,6 @@
+from ..pyutils.cached_property import cached_property
 from ..language import ast
+
 from abc import ABCMeta, abstractmethod
 import six
 
@@ -18,7 +20,12 @@ def __init__(self, schema, document_string, document_ast, execute):
         self.document_ast = document_ast
         self.execute = execute
 
-    def get_operations(self):
+    @cached_property
+    def operations_map(self):
+        '''
+        returns a Mapping of operation names and it's associated types.
+        E.g. {'myQuery': 'query', 'myMutation': 'mutation'}
+        '''
         document_ast = self.document_ast
         operations = {}
         for definition in document_ast.definitions:
@@ -29,3 +36,15 @@ def get_operations(self):
                     operation_name = None
                 operations[operation_name] = definition.operation
         return operations
+
+    def get_operation_type(self, operation_name):
+        '''
+        Returns the operation type ('query', 'mutation', 'subscription' or None)
+        for the given operation_name.
+        If no operation_name is provided (and only one operation exists) it will return the
+        operation type for that operation
+        '''
+        operations_map = self.operations_map
+        if not operation_name and len(operations_map) == 1:
+            return next(iter(operations_map.values()))
+        return operations_map.get(operation_name)

graphql/backend/tests/test_document.py

@@ -0,0 +1,71 @@
+from ...language.base import parse
+from ..base import GraphQLDocument
+from .schema import schema
+
+
+def create_document(document_string):
+    document_ast = parse(document_string)
+    return GraphQLDocument(
+        schema=schema,
+        document_string=document_string,
+        document_ast=document_ast,
+        execute=lambda *_: None
+    )
+
+
+def test_document_operations_map_unnamed_operation():
+    document = create_document('{ hello }')
+    assert document.operations_map == {
+        None: 'query'
+    }
+
+
+def test_document_operations_map_multiple_queries():
+    document = create_document('''
+    query MyQuery1 { hello }
+    query MyQuery2 { hello }
+    ''')
+    assert document.operations_map == {
+        'MyQuery1': 'query',
+        'MyQuery2': 'query'
+    }
+
+
+def test_document_operations_map_multiple_queries():
+    document = create_document('''
+    query MyQuery { hello }
+    mutation MyMutation { hello }
+    subscription MySubscription { hello }
+    ''')
+    assert document.operations_map == {
+        'MyQuery': 'query',
+        'MyMutation': 'mutation',
+        'MySubscription': 'subscription'
+    }
+
+
+def test_document_get_operation_type_unnamed_operation():
+    document = create_document('''
+    query { hello }
+    ''')
+    assert document.get_operation_type(None) == 'query'
+    assert document.get_operation_type('Unknown') is None
+
+
+def test_document_get_operation_type_multiple_operations():
+    document = create_document('''
+    query MyQuery { hello }
+    mutation MyMutation {hello}
+    ''')
+    assert document.get_operation_type(None) is None
+    assert document.get_operation_type('MyQuery') == 'query'
+    assert document.get_operation_type('MyMutation') == 'mutation'
+    assert document.get_operation_type('Unexistent') is None
+
+
+def test_document_get_operation_type_multiple_operations_empty_operation_name():
+    document = create_document('''
+    query MyQuery { hello }
+    mutation {hello}
+    ''')
+    assert document.get_operation_type(None) is 'mutation'