Add feature to flatten nested GraphQL unions. #286

derive/src/args.rs

@@ -447,6 +447,32 @@ impl EnumItem {
     }
 }
 
+pub struct UnionItem {
+    pub flatten: bool,
+}
+
+impl UnionItem {
+    pub fn parse(attrs: &[Attribute]) -> Result<Self> {
+        let mut flatten = false;
+
+        for attr in attrs {
+            if attr.path.is_ident("item") {
+                if let Meta::List(args) = attr.parse_meta()? {
+                    for meta in args.nested {
+                        if let NestedMeta::Meta(Meta::Path(p)) = meta {
+                            if p.is_ident("flatten") {
+                                flatten = true;
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        Ok(Self { flatten })
+    }
+}
+
 pub struct InputField {
     pub name: Option<String>,
     pub desc: Option<String>,

derive/src/lib.rs

@@ -100,7 +100,7 @@ pub fn derive_interface(input: TokenStream) -> TokenStream {
     }
 }
 
-#[proc_macro_derive(Union, attributes(graphql))]
+#[proc_macro_derive(Union, attributes(graphql, item))]
 pub fn derive_union(input: TokenStream) -> TokenStream {
     let (args, input) = match parse_derive(input.into()) {
         Ok(r) => r,

derive/src/union.rs

@@ -58,8 +58,11 @@ pub fn generate(union_args: &args::Interface, input: &DeriveInput) -> Result<Tok
                 ))
             }
         };
+
         if let Type::Path(p) = &field.ty {
             // This validates that the field type wasn't already used
+            let item_args = args::UnionItem::parse(&variant.attrs)?;
+
             if !enum_items.insert(p) {
                 return Err(Error::new_spanned(
                     field,
@@ -68,25 +71,58 @@ pub fn generate(union_args: &args::Interface, input: &DeriveInput) -> Result<Tok
             }
 
             enum_names.push(enum_name);
-            type_into_impls.push(quote! {
-                #crate_name::static_assertions::assert_impl_one!(#p: #crate_name::type_mark::TypeMarkObject);
 
-                #[allow(clippy::all, clippy::pedantic)]
-                impl #generics ::std::convert::From<#p> for #ident #generics {
-                    fn from(obj: #p) -> Self {
-                        #ident::#enum_name(obj)
+            if !item_args.flatten {
+                type_into_impls.push(quote! {
+                    #crate_name::static_assertions::assert_impl_one!(#p: #crate_name::type_mark::TypeMarkObject);
+
+                    #[allow(clippy::all, clippy::pedantic)]
+                    impl #generics ::std::convert::From<#p> for #ident #generics {
+                        fn from(obj: #p) -> Self {
+                            #ident::#enum_name(obj)
+                        }
                     }
-                }
-            });
+                });
+            } else {
+                type_into_impls.push(quote! {
+                    #crate_name::static_assertions::assert_impl_one!(#p: #crate_name::type_mark::TypeMarkEnum);
+
+                    #[allow(clippy::all, clippy::pedantic)]
+                    impl #generics ::std::convert::From<#p> for #ident #generics {
+                        fn from(obj: #p) -> Self {
+                            #ident::#enum_name(obj)
+                        }
+                    }
+                });
+            }
+
             registry_types.push(quote! {
                 <#p as #crate_name::Type>::create_type_info(registry);
             });
-            possible_types.push(quote! {
-                possible_types.insert(<#p as #crate_name::Type>::type_name().to_string());
-            });
-            get_introspection_typename.push(quote! {
-                #ident::#enum_name(obj) => <#p as #crate_name::Type>::type_name()
-            });
+
+            if !item_args.flatten {
+                possible_types.push(quote! {
+                    possible_types.insert(<#p as #crate_name::Type>::type_name().to_string());
+                });
+            } else {
+                possible_types.push(quote! {
+                    if let Some(#crate_name::registry::MetaType::Union { possible_types: possible_types2, .. }) =
+                        registry.types.remove(&*<#p as #crate_name::Type>::type_name()) {
+                        possible_types.extend(possible_types2);
+                    }
+                });
+            }
+
+            if !item_args.flatten {
+                get_introspection_typename.push(quote! {
+                    #ident::#enum_name(obj) => <#p as #crate_name::Type>::type_name()
+                });
+            } else {
+                get_introspection_typename.push(quote! {
+                    #ident::#enum_name(obj) => <#p as #crate_name::Type>::introspection_type_name(obj)
+                });
+            }
+
             collect_all_fields.push(quote! {
                 #ident::#enum_name(obj) => obj.collect_all_fields(ctx, fields)
             });

tests/union.rs

@@ -284,3 +284,74 @@ pub async fn test_union_field_result() {
         })
     );
 }
+
+#[async_std::test]
+pub async fn test_union_flatten() {
+    #[derive(SimpleObject)]
+    struct MyObj1 {
+        value1: i32,
+    }
+
+    #[derive(SimpleObject)]
+    struct MyObj2 {
+        value2: i32,
+    }
+
+    #[derive(Union)]
+    enum InnerUnion1 {
+        A(MyObj1),
+    }
+
+    #[derive(Union)]
+    enum InnerUnion2 {
+        B(MyObj2),
+    }
+
+    #[derive(Union)]
+    enum MyUnion {
+        #[item(flatten)]
+        Inner1(InnerUnion1),
+
+        #[item(flatten)]
+        Inner2(InnerUnion2),
+    }
+
+    struct Query;
+
+    #[Object]
+    impl Query {
+        async fn value1(&self) -> MyUnion {
+            InnerUnion1::A(MyObj1 { value1: 99 }).into()
+        }
+
+        async fn value2(&self) -> MyUnion {
+            InnerUnion2::B(MyObj2 { value2: 88 }).into()
+        }
+    }
+
+    let schema = Schema::new(Query, EmptyMutation, EmptySubscription);
+    let query = r#"
+    {
+        value1 {
+            ... on MyObj1 {
+                value1
+            }
+        }
+        value2 {
+            ... on MyObj2 {
+                value2
+            }
+        }
+    }"#;
+    assert_eq!(
+        schema.execute(query).await.into_result().unwrap().data,
+        serde_json::json!({
+            "value1": {
+                "value1": 99,
+            },
+            "value2": {
+                "value2": 88,
+            }
+        })
+    );
+}