From 391d3d10e435864c082e0846ae22ae8276ec36ff Mon Sep 17 00:00:00 2001
From: Miaxos <an.griffon@gmail.com>
Date: Mon, 11 Oct 2021 12:26:18 +0000
Subject: [PATCH] feat: add visible macro argument for union type

---
 derive/src/args.rs  |  2 ++
 derive/src/union.rs | 16 ++++++++++++++++
 src/lib.rs          |  2 ++
 src/model/type.rs   | 14 +++++++++-----
 src/registry/mod.rs |  8 ++++++++
 5 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/derive/src/args.rs b/derive/src/args.rs
index 070519f3..0cdcf5b2 100644
--- a/derive/src/args.rs
+++ b/derive/src/args.rs
@@ -304,6 +304,8 @@ pub struct UnionItem {
 
     #[darling(default)]
     pub flatten: bool,
+    #[darling(default)]
+    pub visible: Option<Visible>,
 }
 
 #[derive(FromField)]
diff --git a/derive/src/union.rs b/derive/src/union.rs
index e6c1efb5..d1de812a 100644
--- a/derive/src/union.rs
+++ b/derive/src/union.rs
@@ -33,11 +33,13 @@ pub fn generate(union_args: &args::Union) -> GeneratorResult<TokenStream> {
 
     let mut registry_types = Vec::new();
     let mut possible_types = Vec::new();
+    let mut union_values = Vec::new();
     let mut get_introspection_typename = Vec::new();
     let mut collect_all_fields = Vec::new();
 
     for variant in s {
         let enum_name = &variant.ident;
+        let union_visible = visible_fn(&variant.visible);
         let ty = match variant.fields.style {
             Style::Tuple if variant.fields.fields.len() == 1 => &variant.fields.fields[0],
             Style::Tuple => {
@@ -70,6 +72,15 @@ pub fn generate(union_args: &args::Union) -> GeneratorResult<TokenStream> {
             }
 
             enum_names.push(enum_name);
+            union_values.push(quote! {
+                union_values.insert(
+                    <#p as #crate_name::Type>::type_name().into_owned(),
+                    #crate_name::registry::MetaUnionValue {
+                        name: <#p as #crate_name::Type>::type_name().into_owned(),
+                        visible: #union_visible,
+                    }
+                );
+            });
 
             struct RemoveLifetime;
             impl VisitMut for RemoveLifetime {
@@ -176,6 +187,11 @@ pub fn generate(union_args: &args::Union) -> GeneratorResult<TokenStream> {
                             #(#possible_types)*
                             possible_types
                         },
+                        union_values: {
+                            let mut union_values = #crate_name::indexmap::IndexMap::new();
+                            #(#union_values)*
+                            union_values
+                        },
                         visible: #visible,
                     }
                 })
diff --git a/src/lib.rs b/src/lib.rs
index f11916f0..37fa27b5 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -831,6 +831,8 @@ pub use async_graphql_derive::Interface;
 /// | Attribute    | description                              | Type     | Optional |
 /// |--------------|------------------------------------------|----------|----------|
 /// | flatten      | Similar to serde (flatten)               | boolean  | Y        |
+/// | visible       | If `false`, it will not be displayed in introspection. *[See also the Book](https://async-graphql.github.io/async-graphql/en/visibility.html).* | bool | Y |
+/// | visible       | Call the specified function. If the return value is `false`, it will not be displayed in introspection. | string | Y |
 ///
 /// # Define a union
 ///
diff --git a/src/model/type.rs b/src/model/type.rs
index 7741f76b..4f20506c 100644
--- a/src/model/type.rs
+++ b/src/model/type.rs
@@ -133,7 +133,7 @@ impl<'a> __Type<'a> {
         }
     }
 
-    async fn possible_types(&self) -> Option<Vec<__Type<'a>>> {
+    async fn possible_types(&self, ctx: &Context<'_>) -> Option<Vec<__Type<'a>>> {
         if let TypeDetail::Named(registry::MetaType::Interface { possible_types, .. }) =
             &self.detail
         {
@@ -143,13 +143,17 @@ impl<'a> __Type<'a> {
                     .map(|ty| __Type::new(self.registry, ty))
                     .collect(),
             )
-        } else if let TypeDetail::Named(registry::MetaType::Union { possible_types, .. }) =
+        } else if let TypeDetail::Named(registry::MetaType::Union { union_values, .. }) =
             &self.detail
         {
             Some(
-                possible_types
-                    .iter()
-                    .map(|ty| __Type::new(self.registry, ty))
+                union_values
+                    .values()
+                    .filter(|value| match &value.visible {
+                        Some(f) => f(ctx),
+                        None => true,
+                    })
+                    .map(|ty| __Type::new(self.registry, &ty.name))
                     .collect(),
             )
         } else {
diff --git a/src/registry/mod.rs b/src/registry/mod.rs
index 041fb936..254cae2c 100644
--- a/src/registry/mod.rs
+++ b/src/registry/mod.rs
@@ -174,6 +174,12 @@ pub struct MetaEnumValue {
     pub visible: Option<MetaVisibleFn>,
 }
 
+#[derive(Clone)]
+pub struct MetaUnionValue {
+    pub name: String,
+    pub visible: Option<MetaVisibleFn>,
+}
+
 type MetaVisibleFn = fn(&Context<'_>) -> bool;
 
 #[derive(Clone)]
@@ -205,6 +211,7 @@ pub enum MetaType {
     Union {
         name: String,
         description: Option<&'static str>,
+        union_values: IndexMap<String, MetaUnionValue>,
         possible_types: IndexSet<String>,
         visible: Option<MetaVisibleFn>,
     },
@@ -461,6 +468,7 @@ impl Registry {
             MetaType::Union {
                 name: "_Entity".to_string(),
                 description: None,
+                union_values: Default::default(),
                 possible_types,
                 visible: None,
             },