WIP union selection fix

tomhoule · tomhoule · commit e565732407a1 · 2019-01-05T10:47:11.000+01:00
diff --git a/graphql_client/tests/interfaces/interface_with_fragment_query.graphql b/graphql_client/tests/interfaces/interface_with_fragment_query.graphql
@@ -1,4 +1,5 @@
 fragment PublicStatus on Named {
+  __typename
   displayName
 }
 
diff --git a/graphql_client_codegen/src/fragments.rs b/graphql_client_codegen/src/fragments.rs
@@ -1,4 +1,4 @@
-use proc_macro2::{Ident, Span, TokenStream};
+use proc_macro2::TokenStream;
 use query::QueryContext;
 use selection::Selection;
 use std::cell::Cell;
@@ -19,35 +19,16 @@ pub(crate) struct GqlFragment {
 impl GqlFragment {
     /// Generate all the Rust code required by the fragment's selection.
     pub(crate) fn to_rust(&self, context: &QueryContext) -> Result<TokenStream, ::failure::Error> {
-        let derives = context.response_derives();
-        let name_ident = Ident::new(&self.name, Span::call_site());
-        let opt_object = context.schema.objects.get(&self.on);
-        let (field_impls, fields) = if let Some(object) = opt_object {
-            let field_impls =
-                object.field_impls_for_selection(context, &self.selection, &self.name)?;
-            let fields =
-                object.response_fields_for_selection(context, &self.selection, &self.name)?;
-            (field_impls, fields)
+        if let Some(obj) = context.schema.objects.get(&self.on) {
+            obj.response_for_selection(context, &self.selection, &self.name)
         } else if let Some(iface) = context.schema.interfaces.get(&self.on) {
-            let field_impls =
-                iface.field_impls_for_selection(context, &self.selection, &self.name)?;
-            let fields =
-                iface.response_fields_for_selection(context, &self.selection, &self.name)?;
-            (field_impls, fields)
+            iface.response_for_selection(context, &self.selection, &self.name)
         } else {
-            panic!(
-                "fragment '{}' cannot operate on unknown type '{}'",
-                self.name, self.on
-            );
-        };
-
-        Ok(quote! {
-            #derives
-            pub struct #name_ident {
-                #(#fields,)*
-            }
-
-            #(#field_impls)*
-        })
+            Err(format_err!(
+                "Fragment {} is defined on unknown type: {}",
+                self.name,
+                self.on
+            ))?
+        }
     }
 }
diff --git a/graphql_client_codegen/src/query.rs b/graphql_client_codegen/src/query.rs
@@ -48,6 +48,7 @@ impl QueryContext {
         }
     }
 
+    ///
     pub(crate) fn maybe_expand_field(
         &self,
         ty: &str,
@@ -67,7 +68,7 @@ impl QueryContext {
             unn.is_required.set(true);
             unn.response_for_selection(self, &selection, prefix)
         } else {
-            Ok(quote!())
+            Err(format_err!("Unknown type: {}", ty))
         }
     }
 
diff --git a/graphql_client_codegen/src/selection.rs b/graphql_client_codegen/src/selection.rs
@@ -1,29 +1,31 @@
 use constants::*;
 use graphql_parser::query::SelectionSet;
+use std::collections::BTreeMap;
+use std::error::Error;
 
 /// A single object field as part of a selection.
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct SelectionField {
     pub alias: Option<String>,
     pub name: String,
     pub fields: Selection,
 }
 
 /// A spread fragment in a selection (e.g. `...MyFragment`).
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct SelectionFragmentSpread {
     pub fragment_name: String,
 }
 
 /// An inline fragment as part of a selection (e.g. `...on MyThing { name }`).
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct SelectionInlineFragment {
     pub on: String,
     pub fields: Selection,
 }
 
 /// An element in a query selection.
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub enum SelectionItem {
     Field(SelectionField),
     FragmentSpread(SelectionFragmentSpread),
@@ -41,7 +43,7 @@ impl SelectionItem {
     }
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct Selection(pub Vec<SelectionItem>);
 
 impl Selection {
@@ -71,6 +73,57 @@ impl Selection {
             .next()
     }
 
+    // Implementation helper for `selected_variants_on_union`.
+    fn selected_variants_on_union_inner<'s>(
+        &'s self,
+        context: &'s crate::query::QueryContext,
+        selected_variants: &mut BTreeMap<&'s str, Vec<&'s SelectionItem>>,
+    ) -> Result<(), Box<Error>> {
+        for item in self.0.iter() {
+            match item {
+                SelectionItem::Field(_) => Err(format_err!("field selection on union"))?,
+                SelectionItem::InlineFragment(inline_fragment) => {
+                    selected_variants
+                        .entry(inline_fragment.on.as_str())
+                        .and_modify(|entry| entry.extend(&inline_fragment.fields.0))
+                        .or_insert_with(|| {
+                            let mut set = Vec::new();
+                            set.extend(&inline_fragment.fields.0);
+                            set
+                        });
+                }
+                SelectionItem::FragmentSpread(SelectionFragmentSpread { fragment_name }) => {
+                    let fragment = context
+                        .fragments
+                        .get(fragment_name)
+                        .ok_or_else(|| format_err!("Unknown fragment: {}", &fragment_name))?;
+
+                    fragment
+                        .selection
+                        .selected_variants_on_union_inner(context, selected_variants)?;
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    /// This method should only be invoked on selections on union fields. It returns the selected varianst and the
+    ///
+    /// Importantly, it will "flatten" the fragments and handle multiple selections of the same variant.
+    ///
+    /// The `context` argument is required so we can expand the fragments.
+    pub(crate) fn selected_variants_on_union<'s>(
+        &'s self,
+        context: &'s crate::query::QueryContext,
+    ) -> Result<BTreeMap<&'s str, Vec<&'s SelectionItem>>, Box<Error>> {
+        let mut selected_variants = BTreeMap::new();
+
+        self.selected_variants_on_union_inner(context, &mut selected_variants)?;
+
+        Ok(selected_variants)
+    }
+
     #[cfg(test)]
     pub(crate) fn new_empty() -> Selection {
         Selection(Vec::new())
diff --git a/graphql_client_codegen/src/unions.rs b/graphql_client_codegen/src/unions.rs
@@ -7,7 +7,7 @@ use std::cell::Cell;
 use std::collections::BTreeSet;
 
 #[derive(Debug, Clone, PartialEq)]
-pub struct GqlUnion {
+pub(crate) struct GqlUnion {
     pub description: Option<String>,
     pub variants: BTreeSet<String>,
     pub is_required: Cell<bool>,
@@ -22,90 +22,104 @@ enum UnionError {
     MissingTypename { union_name: String },
 }
 
-type UnionVariantResult = Result<(Vec<TokenStream>, Vec<TokenStream>, Vec<String>), failure::Error>;
-
-pub(crate) fn union_variants(
-    selection: &Selection,
-    query_context: &QueryContext,
+type UnionVariantResult<'selection> =
+    Result<(Vec<TokenStream>, Vec<TokenStream>, Vec<&'selection str>), failure::Error>;
+
+/// Returns a triple.
+///
+/// - The first element is the union variants to be inserted directly into the `enum` declaration.
+/// - The second is the structs for each variant's sub-selection
+/// - The last one contains which fields have been selected on the union, so we can make the enum exhaustive by complementing with those missing.
+pub(crate) fn union_variants<'selection>(
+    selection: &'selection Selection,
+    context: &QueryContext,
     prefix: &str,
-) -> UnionVariantResult {
-    let mut children_definitions = Vec::new();
-    let mut used_variants = Vec::with_capacity(selection.0.len());
-
-    let variants: Result<Vec<TokenStream>, failure::Error> = selection
-        .0
-        .iter()
-        // ignore __typename
-        .filter(|item| {
-            if let SelectionItem::Field(f) = item {
-                f.name != TYPENAME_FIELD
-            } else {
-                true
-            }
-        })
-        .map(|item| {
-            let (on, fields) = match item {
-                SelectionItem::Field(_) => Err(format_err!("field selection on union"))?,
-                SelectionItem::FragmentSpread(SelectionFragmentSpread { fragment_name }) => {
-                    let fragment = query_context
-                        .fragments
-                        .get(fragment_name)
-                        .ok_or_else(|| format_err!("Unknown fragment: {}", &fragment_name))?;
-
-                    (&fragment.on, &fragment.selection)
-                }
-                SelectionItem::InlineFragment(frag) => (&frag.on, &frag.fields),
-            };
-            let variant_name = Ident::new(&on, Span::call_site());
-            used_variants.push(on.to_string());
-
-            let new_prefix = format!("{}On{}", prefix, on);
-
-            let variant_type = Ident::new(&new_prefix, Span::call_site());
-
-            let field_object_type = query_context
-                .schema
-                .objects
-                .get(on)
-                .map(|_f| query_context.maybe_expand_field(&on, &fields, &new_prefix));
-            let field_interface = query_context
-                .schema
-                .interfaces
-                .get(on)
-                .map(|_f| query_context.maybe_expand_field(&on, &fields, &new_prefix));
-            let field_union_type = query_context
-                .schema
-                .unions
-                .get(on)
-                .map(|_f| query_context.maybe_expand_field(&on, &fields, &new_prefix));
-
-            match field_object_type.or(field_interface).or(field_union_type) {
-                Some(tokens) => children_definitions.push(tokens?),
-                None => Err(UnionError::UnknownType { ty: on.to_string() })?,
-            };
-
-            Ok(quote! {
-                #variant_name(#variant_type)
-            })
-        })
-        .collect();
+) -> UnionVariantResult<'selection> {
+    let selection = selection.selected_variants_on_union(context)?;
+    let used_variants = selection.keys().collect();
+    let mut children_definitions = Vec::with_capacity(selection.size());
+    let mut variants = Vec::with_capacity(selection.size());
+
+    for (on, fields) in selection.iter() {
+        let variant_name = Ident::new(&on, Span::call_site());
+        used_variants.push(on.to_string());
+
+        let new_prefix = format!("{}On{}", prefix, on);
+
+        let variant_type = Ident::new(&new_prefix, Span::call_site());
+
+        let field_object_type = context
+            .schema
+            .objects
+            .get(on)
+            .map(|_f| context.maybe_expand_field(&on, fields, &new_prefix));
+        let field_interface = context
+            .schema
+            .interfaces
+            .get(on)
+            .map(|_f| context.maybe_expand_field(&on, fields, &new_prefix));
+        let field_union_type = context
+            .schema
+            .unions
+            .get(on)
+            .map(|_f| context.maybe_expand_field(&on, fields, &new_prefix));
+
+        match field_object_type.or(field_interface).or(field_union_type) {
+            Some(tokens) => children_definitions.push(tokens?),
+            None => Err(UnionError::UnknownType { ty: on.to_string() })?,
+        };
 
-    let variants = variants?;
+        variants.push(quote! {
+            #variant_name(#variant_type)
+        })
+    }
 
     Ok((variants, children_definitions, used_variants))
+
+    // let variants: Result<Vec<TokenStream>, failure::Error> = selection
+    //     .0
+    //     .iter()
+    //     // ignore __typename
+    //     .filter(|item| {
+    //         if let SelectionItem::Field(f) = item {
+    //             f.name != TYPENAME_FIELD
+    //         } else {
+    //             true
+    //         }
+    //     })
+    //     // .flat_map(
+    //     //     |item| -> impl ::std::iter::Iterator<Item = Result<(_, _), failure::Error>> {
+    //     //         match item {
+    //     //             SelectionItem::Field(_) => Err(format_err!("field selection on union"))?,
+    //     //             SelectionItem::FragmentSpread(SelectionFragmentSpread { fragment_name }) => {
+    //     //                 let fragment = query_context
+    //     //                     .fragments
+    //     //                     .get(fragment_name)
+    //     //                     .ok_or_else(|| format_err!("Unknown fragment: {}", &fragment_name))?;
+    //     //                 // found the bug! `on` doesn't mean the same here as in the inline fragments
+    //     //                 std::iter::once(Ok((&fragment.on, &fragment.selection)))
+    //     //             }
+    //     //             SelectionItem::InlineFragment(frag) => {
+    //     //                 std::iter::once(Ok((&frag.on, &frag.fields)))
+    //     //             }
+    //     //         }
+    //     //     },
+    //     // )
+    //     // // .collect::<Result<_, _>>()?
+    //     .map(|result: Result<(_, _), failure::Error>| -> Result<_, _> {
+    //         let Ok((on, fields)) = result?;
+
+    // let variants = variants?;
 }
 
 impl GqlUnion {
+    /// Returns the code to deserialize this union in the response given the query selection.
     pub(crate) fn response_for_selection(
         &self,
         query_context: &QueryContext,
         selection: &Selection,
         prefix: &str,
     ) -> Result<TokenStream, failure::Error> {
-        let struct_name = Ident::new(prefix, Span::call_site());
-        let derives = query_context.response_derives();
-
-        // TODO: do this inside fragments
         let typename_field = selection.extract_typename(query_context);
 
         if typename_field.is_none() {
@@ -114,6 +128,9 @@ impl GqlUnion {
             })?;
         }
 
+        let struct_name = Ident::new(prefix, Span::call_site());
+        let derives = query_context.response_derives();
+
         let (mut variants, children_definitions, used_variants) =
             union_variants(selection, query_context, prefix)?;
 

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`fragment PublicStatus on Named {`
	`2`	`+ __typename`
`2`	`3`	`displayName`
`3`	`4`	`}`
`4`	`5`
Original file line number	Diff line number	Diff line change
`@@ -48,6 +48,7 @@ impl QueryContext {`
`48`	`48`	`}`
`49`	`49`	`}`
`50`	`50`
	`51`	`+ ///`
`51`	`52`	`pub(crate) fn maybe_expand_field(`
`52`	`53`	`&self,`
`53`	`54`	`ty: &str,`
`@@ -67,7 +68,7 @@ impl QueryContext {`
`67`	`68`	`unn.is_required.set(true);`
`68`	`69`	`unn.response_for_selection(self, &selection, prefix)`
`69`	`70`	`} else {`
`70`		`- Ok(quote!())`
	`71`	`+ Err(format_err!("Unknown type: {}", ty))`
`71`	`72`	`}`
`72`	`73`	`}`
`73`	`74`