/
helpers.py
185 lines (144 loc) · 7.12 KB
/
helpers.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
# Copyright 2019-present Kensho Technologies, LLC.
from copy import copy
from graphql.language.ast import Field, InlineFragment, OperationDefinition, SelectionSet
from ..macro_edge.directives import MacroEdgeTargetDirective
def _yield_ast_nodes_with_directives(ast):
"""Get the AST objects where directives appear, anywhere in the given AST.
Args:
ast: GraphQL library AST object, such as a Field, InlineFragment, or OperationDefinition
Returns:
Iterable[Tuple[AST object, Directive]], where each tuple describes an AST node together with
the directive it contains. If an AST node contains multiple directives, the AST node will be
returned as part of multiple tuples, in no particular order.
"""
for directive in ast.directives:
yield (ast, directive)
if isinstance(ast, (Field, InlineFragment, OperationDefinition)):
if ast.selection_set is not None:
for sub_selection_set in ast.selection_set.selections:
# TODO(predrag): When we make the compiler py3-only, use a "yield from" here.
for entry in _yield_ast_nodes_with_directives(sub_selection_set):
yield entry
else:
raise AssertionError(u'Unexpected AST type received: {} {}'.format(type(ast), ast))
def get_directives_for_ast(ast):
"""Return a dict of directive name -> list of (ast, directive) where that directive is used.
Args:
ast: GraphQL library AST object, such as a Field, InlineFragment, or OperationDefinition
Returns:
Dict[str, List[Tuple[AST object, Directive]]], allowing the user to find the instances
in this AST object where a directive with a given name appears; for each of those instances,
we record and return the AST object where the directive was applied, together with the AST
Directive object describing it together with any arguments that might have been supplied.
"""
result = {}
for ast, directive in _yield_ast_nodes_with_directives(ast):
directive_name = directive.name.value
result.setdefault(directive_name, []).append((ast, directive))
return result
def remove_directives_from_ast(ast, directive_names_to_omit):
"""Return an equivalent AST to the input, 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
def find_target_and_copy_path_to_it(ast):
"""Copy the AST objects on the path to the target, returning it and the target AST.
This function makes it easy to change the AST at the @target directive without mutating the
original object or doing a deepcopy.
Args:
ast: GraphQL library AST object
Returns:
pair containing:
- GraphQL library AST object equal to the input. Objects on the path to the @target
directive are shallow copied.
- GraphQL library AST object at the @target directive of the resulting AST, or None
if there was no @target directive in the AST.
"""
for directive in ast.directives:
if directive.name.value == MacroEdgeTargetDirective.name:
target_ast = copy(ast)
return target_ast, target_ast
new_selections = []
target_ast = None
if isinstance(ast, (Field, InlineFragment, OperationDefinition)):
if ast.selection_set is not None:
for selection in ast.selection_set.selections:
new_selection, possible_target_ast = find_target_and_copy_path_to_it(selection)
new_selections.append(new_selection)
if possible_target_ast is not None:
target_ast = possible_target_ast
else:
raise AssertionError(u'Unexpected AST type received: {} {}'.format(type(ast), ast))
if target_ast is None:
return ast, None
else:
new_ast = copy(ast)
new_ast.selection_set = SelectionSet(new_selections)
return new_ast, target_ast
def omit_ast_from_ast_selections(ast, ast_to_omit):
"""Return an equivalent AST to the input, but with the specified AST omitted if it appears.
Args:
ast: GraphQL library AST object, such as a Field, InlineFragment, or OperationDefinition
ast_to_omit: GraphQL library AST object, the *exact same* object that should be omitted.
This function uses reference equality, since deep equality can get expensive.
Returns:
GraphQL library AST object, equivalent to the input one, with all instances of
the specified AST omitted. If the specified AST does 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
if ast.selection_set is None:
return ast
made_changes = False
selections_to_keep = []
for selection_ast in ast.selection_set.selections:
if selection_ast is ast_to_omit:
# Drop the current selection.
made_changes = True
else:
new_selection_ast = omit_ast_from_ast_selections(selection_ast, ast_to_omit)
if new_selection_ast is not selection_ast:
# The current selection contained the AST to omit, and was altered as a result.
made_changes = True
selections_to_keep.append(new_selection_ast)
if not made_changes:
return ast
new_ast = copy(ast)
if not selections_to_keep:
new_ast.selection_set = None
else:
new_ast.selection_set = SelectionSet(selections_to_keep)
return new_ast