Implement get_operation_ast utility

#5

Author: jhgg

graphql/core/utils/get_operation_ast.py

@@ -0,0 +1,21 @@
+from ...core.language import ast
+
+
+def get_operation_ast(document_ast, operation_name=None):
+    operation = None
+
+    for definition in document_ast.definitions:
+        if isinstance(definition, ast.OperationDefinition):
+            if not operation_name:
+                # If no operation name is provided, only return an Operation if it is the only one present in the
+                # document. This means that if we've encountered a second operation as we were iterating over the
+                # definitions in the document, there are more than one Operation defined, and we should return None.
+                if operation:
+                    return None
+
+                operation = definition
+
+            elif definition.name and definition.name.value == operation_name:
+                return definition
+
+    return operation

tests/core_utils/test_get_operation_ast.py

@@ -0,0 +1,46 @@
+from graphql.core import parse
+from graphql.core.utils.get_operation_ast import get_operation_ast
+
+
+def test_gets_an_operation_from_a_simple_document():
+    doc = parse('{ field }')
+    assert get_operation_ast(doc) == doc.definitions[0]
+
+
+def test_gets_an_operation_from_a_document_with_named_operation():
+    doc = parse('mutation Test { field }')
+    assert get_operation_ast(doc) == doc.definitions[0]
+
+
+def test_does_not_get_missing_operation():
+    doc = parse('{ field } mutation Test { field }')
+    assert not get_operation_ast(doc)
+
+
+def test_does_not_get_ambiguous_unnamed_operation():
+    doc = parse('query TestQ { field } mutation TestM { field }')
+    assert not get_operation_ast(doc)
+
+
+def test_does_not_get_misnamed_operation():
+    doc = parse('query TestQ { field } mutation TestM { field }')
+    assert not get_operation_ast(doc, 'Unknown')
+
+
+def test_gets_named_operation():
+    doc = parse('query TestQ { field } mutation TestM { field }')
+    assert get_operation_ast(doc, 'TestQ') == doc.definitions[0]
+    assert get_operation_ast(doc, 'TestM') == doc.definitions[1]
+
+
+def test_does_not_get_fragment():
+    doc = parse('fragment Foo on Type { field }')
+    assert not get_operation_ast(doc)
+    assert not get_operation_ast(doc, 'Foo')
+
+
+def test_does_not_get_fragment_with_same_name_query():
+    doc = parse('fragment Foo on Type { field } query Foo { field }')
+    assert get_operation_ast(doc) == doc.definitions[1]
+    assert get_operation_ast(doc, 'Foo') == doc.definitions[1]
+