Add more helpers, rename the file.

graphql_compiler/macros/macro_edge/validation_helpers.py → graphql_compiler/macros/macro_edge/helpers.py

@@ -1,5 +1,9 @@
 # Copyright 2019-present Kensho Technologies, LLC.
-from graphql.language.ast import Field, InlineFragment, OperationDefinition
+from copy import copy
+
+from graphql.language.ast import (
+    Argument, Directive, Field, InlineFragment, Name, OperationDefinition, SelectionSet, StringValue
+)
 
 from ...ast_manipulation import get_human_friendly_ast_field_name
 from ...exceptions import GraphQLInvalidMacroError
@@ -70,3 +74,51 @@ def get_only_selection_from_ast(ast):
                                        .format(len(selection_names), selection_names, ast_name))
 
     return selections[0]
+
+
+def remove_directives_from_ast(ast, directive_names_to_omit):
+    """Return a copy of the AST but with instances of the named directives omitted.
+
+    Args:
+        ast: GraphQL library AST object, such as a Field, InlineFragment, or OperationDefinition
+        directive_names_to_omit: set of strings describing the names of the directives to omit
+
+    Returns:
+        GraphQL library AST object, equivalent to the input one, with all instances of
+        the named directives omitted. If the specified directives do not appear in the input AST,
+        the returned object is the exact same object as the input.
+    """
+    if not isinstance(ast, (Field, InlineFragment, OperationDefinition)):
+        return ast
+
+    made_changes = False
+
+    new_selections = None
+    if ast.selection_set is not None:
+        new_selections = []
+        for selection_ast in ast.selection_set.selections:
+            new_selection_ast = remove_directives_from_ast(selection_ast, directive_names_to_omit)
+
+            if selection_ast is not new_selection_ast:
+                # Since we did not get the exact same object as the input, changes were made.
+                # That means this call will also need to make changes and return a new object.
+                made_changes = True
+
+            new_selections.append(new_selection_ast)
+
+    directives_to_keep = [
+        directive
+        for directive in ast.directives
+        if directive.name.value not in directive_names_to_omit
+    ]
+    if len(directives_to_keep) != len(ast.directives):
+        made_changes = True
+
+    if not made_changes:
+        # We didn't change anything, return the original input object.
+        return ast
+
+    new_ast = copy(ast)
+    new_ast.selection_set = SelectionSet(new_selections)
+    new_ast.directives = directives_to_keep
+    return new_ast