Add output info to query metadata (kensho-technologies#311)

* Add output info to query metadata

* Update explain info tests

* Add output tests to explain_info

* Remove context output

* Add docstring metadata outputs property

* Use query_metadta_table variable and name

* Replace more names with query_metadata_table

* Remove fold from OutputInfo

Author: evantey14

graphql_compiler/compiler/compiler_frontend.py

@@ -90,7 +90,7 @@
     get_vertex_field_type, invert_dict, is_tagged_parameter, is_vertex_field_name,
     strip_non_null_from_type, validate_output_name, validate_safe_string
 )
-from .metadata import LocationInfo, QueryMetadataTable, RecurseInfo, TagInfo
+from .metadata import LocationInfo, OutputInfo, QueryMetadataTable, RecurseInfo, TagInfo
 
 
 # LocationStackEntry contains the following:
@@ -298,7 +298,7 @@ def _compile_property_ast(schema, current_schema_type, ast, location,
     if output_directive:
         # Schema validation has ensured that the fields below exist.
         output_name = output_directive.arguments[0].value.value
-        if output_name in context['outputs']:
+        if context['metadata'].get_output_info(output_name):
             raise GraphQLCompilationError(u'Cannot reuse output name: '
                                           u'{}, {}'.format(output_name, context))
         validate_safe_string(output_name)
@@ -312,12 +312,12 @@ def _compile_property_ast(schema, current_schema_type, ast, location,
             if location.field != COUNT_META_FIELD_NAME:
                 graphql_type = GraphQLList(graphql_type)
 
-        context['outputs'][output_name] = {
-            'location': location,
-            'optional': is_in_optional_scope(context),
-            'type': graphql_type,
-            'fold': context.get('fold', None),
-        }
+        output_info = OutputInfo(
+            location=location,
+            type=graphql_type,
+            optional=is_in_optional_scope(context),
+        )
+        context['metadata'].record_output_info(output_name, output_info)
 
 
 def _get_recurse_directive_depth(field_name, field_directives):
@@ -524,7 +524,7 @@ def _compile_vertex_ast(schema, current_schema_type, ast,
         if edge_traversal_is_folded:
             has_count_filter = has_fold_count_filter(context)
             _validate_fold_has_outputs_or_count_filter(
-                get_context_fold_info(context), has_count_filter, context['outputs'])
+                get_context_fold_info(context), has_count_filter, query_metadata_table)
             basic_blocks.append(blocks.Unfold())
             unmark_context_fold_scope(context)
             if has_count_filter:
@@ -560,7 +560,19 @@ def _compile_vertex_ast(schema, current_schema_type, ast,
     return basic_blocks
 
 
-def _validate_fold_has_outputs_or_count_filter(fold_scope_location, fold_has_count_filter, outputs):
+def _are_locations_in_same_fold(first_location, second_location):
+    """Returns True if locations are contained in the same fold scope."""
+    return (
+        isinstance(first_location, FoldScopeLocation) and
+        isinstance(second_location, FoldScopeLocation) and
+        first_location.base_location == second_location.base_location and
+        first_location.get_first_folded_edge() == second_location.get_first_folded_edge()
+    )
+
+
+def _validate_fold_has_outputs_or_count_filter(
+    fold_scope_location, fold_has_count_filter, query_metadata_table
+):
     """Ensure the @fold scope has at least one output, or filters on the size of the fold."""
     # This function makes sure that the @fold scope has an effect.
     # Folds either output data, or filter the data enclosing the fold based on the size of the fold.
@@ -570,8 +582,8 @@ def _validate_fold_has_outputs_or_count_filter(fold_scope_location, fold_has_cou
 
     # At least one output in the outputs list must point to the fold_scope_location,
     # or the scope corresponding to fold_scope_location had no @outputs and is illegal.
-    for output in six.itervalues(outputs):
-        if output['fold'] == fold_scope_location:
+    for _, output_info in query_metadata_table.outputs:
+        if _are_locations_in_same_fold(output_info.location, fold_scope_location):
             return True
 
     raise GraphQLCompilationError(u'Found a @fold scope that has no effect on the query. '
@@ -798,13 +810,6 @@ def _compile_root_ast_to_ir(schema, ast, type_equivalence_hints=None):
         # query processing, but apply to the global query scope and should be appended to the
         # IR blocks only after the GlobalOperationsStart block has been emitted.
         'global_filters': [],
-        # 'outputs' is a dict mapping each output name to another dict which contains
-        #  - location: Location where to output from
-        #  - optional: boolean representing whether the output was defined within an @optional scope
-        #  - type: GraphQLType of the output
-        #  - fold: FoldScopeLocation object if the current output was defined within a fold scope,
-        #          and None otherwise
-        'outputs': dict(),
         # 'inputs' is a dict mapping input parameter names to their respective expected GraphQL
         # types, as automatically inferred by inspecting the query structure
         'inputs': dict(),
@@ -853,11 +858,10 @@ def _compile_root_ast_to_ir(schema, ast, type_equivalence_hints=None):
     basic_blocks.extend(context['global_filters'])
 
     # Based on the outputs context data, add an output step and construct the output metadata.
-    outputs_context = context['outputs']
-    basic_blocks.append(_compile_output_step(outputs_context))
+    basic_blocks.append(_compile_output_step(query_metadata_table))
     output_metadata = {
-        name: OutputMetadata(type=value['type'], optional=value['optional'])
-        for name, value in six.iteritems(outputs_context)
+        name: OutputMetadata(type=info.type, optional=info.optional)
+        for name, info in query_metadata_table.outputs
     }
 
     return IrAndMetadata(
@@ -867,34 +871,35 @@ def _compile_root_ast_to_ir(schema, ast, type_equivalence_hints=None):
         query_metadata_table=context['metadata'])
 
 
-def _compile_output_step(outputs):
+def _compile_output_step(query_metadata_table):
     """Construct the final ConstructResult basic block that defines the output format of the query.
 
     Args:
-        outputs: dict, output name (string) -> output data dict, specifying the location
-                 from where to get the data, and whether the data is optional (and therefore
-                 may be missing); missing optional data is replaced with 'null'
+        query_metadata_table: QueryMetadataTable object, part of which specifies the location from
+                              where to get the output, and whether the output is optional (and
+                              therefore may be missing); missing optional data is replaced with
+                              'null'
 
     Returns:
         a ConstructResult basic block that constructs appropriate outputs for the query
     """
-    if not outputs:
+    if next(query_metadata_table.outputs, None) is None:
         raise GraphQLCompilationError(u'No fields were selected for output! Please mark at least '
                                       u'one field with the @output directive.')
 
     output_fields = {}
-    for output_name, output_context in six.iteritems(outputs):
-        location = output_context['location']
-        optional = output_context['optional']
-        graphql_type = output_context['type']
+    for output_name, output_info in query_metadata_table.outputs:
+        location = output_info.location
+        optional = output_info.optional
+        graphql_type = output_info.type
 
         expression = None
         existence_check = None
         # pylint: disable=redefined-variable-type
         if isinstance(location, FoldScopeLocation):
             if optional:
                 raise AssertionError(u'Unreachable state reached, optional in fold: '
-                                     u'{}'.format(output_context))
+                                     u'{}'.format(output_info))
 
             if location.field == COUNT_META_FIELD_NAME:
                 expression = expressions.FoldCountContextField(location)

graphql_compiler/compiler/metadata.py

@@ -24,6 +24,15 @@
 )
 
 
+OutputInfo = namedtuple(
+    'OutputInfo',
+    (
+        'location',     # Location/FoldScopeLocation, where to output from
+        'type',         # GraphQLType of the output
+        'optional',     # boolean, whether the output was defined within an @optional scope
+    )
+)
+
 TagInfo = namedtuple(
     'TagInfo',
     (
@@ -161,6 +170,25 @@ def get_location_info(self, location):
                                  u'{}'.format(location))
         return location_info
 
+    @property
+    def outputs(self):
+        """Return an iterable of (output_name, output_info) tuples for all outputs in the query."""
+        for output_name, output_info in six.iteritems(self._outputs):
+            yield output_name, output_info
+
+    def record_output_info(self, output_name, output_info):
+        """Record information about the output."""
+        old_info = self._outputs.get(output_name, None)
+        if old_info is not None:
+            raise AssertionError(u'Attempting to reuse an already-defined output name {}. '
+                                 u'old info {}, new info {}.'
+                                 .format(output_name, old_info, output_info))
+        self._outputs[output_name] = output_info
+
+    def get_output_info(self, output_name):
+        """Get information about an output."""
+        return self._outputs.get(output_name, None)
+
     @property
     def tags(self):
         """Return an iterable of (tag_name, tag_info) tuples for all tags in the query."""