From 49929eae229a2276d1e18f94570c635520e3396e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kat=20March=C3=A1n?= <kzm@zkat.tech>
Date: Sun, 17 Apr 2022 19:30:42 -0700
Subject: [PATCH] feat(labels): allow optional labels in derive macro

Fixes: https://github.com/zkat/miette/issues/110
---
 miette-derive/src/label.rs | 125 ++++++++++++++++++++++++++++---------
 src/handlers/graphical.rs  |   4 +-
 src/lib.rs                 |   8 +++
 tests/derive.rs            |  12 ++++
 4 files changed, 117 insertions(+), 32 deletions(-)

diff --git a/miette-derive/src/label.rs b/miette-derive/src/label.rs
index 5aed762..1d45d7b 100644
--- a/miette-derive/src/label.rs
+++ b/miette-derive/src/label.rs
@@ -18,20 +18,35 @@ pub struct Labels(Vec<Label>);
 
 struct Label {
     label: Option<Display>,
+    optional: bool,
     span: syn::Member,
 }
 
 struct LabelAttr {
     label: Option<Display>,
+    optional: bool,
 }
 
 impl Parse for LabelAttr {
     fn parse(input: ParseStream) -> syn::Result<Self> {
         let la = input.lookahead1();
-        let label = if la.peek(syn::token::Paren) {
+        let (label, optional) = if la.peek(syn::token::Paren) {
             // #[label("{}", x)]
             let content;
             parenthesized!(content in input);
+            let optional = if content.peek(syn::Ident) {
+                let ident = content.parse::<syn::Ident>()?;
+                if ident == "optional" {
+                    if content.peek(syn::Token![,]) {
+                        content.parse::<syn::Token![,]>()?;
+                    }
+                    true
+                } else {
+                    false
+                }
+            } else {
+                false
+            };
             if content.peek(syn::LitStr) {
                 let fmt = content.parse()?;
                 let args = if content.is_empty() {
@@ -44,22 +59,27 @@ impl Parse for LabelAttr {
                     args,
                     has_bonus_display: false,
                 };
-                Some(display)
+                (Some(display), optional)
+            } else if optional {
+                (None, optional)
             } else {
-                return Err(syn::Error::new(input.span(), "Invalid argument to label() attribute. The first argument must be a literal string."));
+                return Err(syn::Error::new(input.span(), "Invalid argument to label() attribute. The first argument must be a literal string or the identifier `optional`"));
             }
         } else if la.peek(Token![=]) {
             // #[label = "blabla"]
             input.parse::<Token![=]>()?;
-            Some(Display {
-                fmt: input.parse()?,
-                args: TokenStream::new(),
-                has_bonus_display: false,
-            })
+            (
+                Some(Display {
+                    fmt: input.parse()?,
+                    args: TokenStream::new(),
+                    has_bonus_display: false,
+                }),
+                false,
+            )
         } else {
-            None
+            (None, false)
         };
-        Ok(LabelAttr { label })
+        Ok(LabelAttr { label, optional })
     }
 }
 
@@ -87,8 +107,13 @@ impl Labels {
                             span: field.span(),
                         })
                     };
-                    let LabelAttr { label } = syn::parse2::<LabelAttr>(attr.tokens.clone())?;
-                    labels.push(Label { label, span });
+                    let LabelAttr { label, optional } =
+                        syn::parse2::<LabelAttr>(attr.tokens.clone())?;
+                    labels.push(Label {
+                        label,
+                        span,
+                        optional,
+                    });
                 }
             }
         }
@@ -102,21 +127,44 @@ impl Labels {
     pub(crate) fn gen_struct(&self, fields: &syn::Fields) -> Option<TokenStream> {
         let (display_pat, display_members) = display_pat_members(fields);
         let labels = self.0.iter().map(|highlight| {
-            let Label { span, label } = highlight;
+            let Label {
+                span,
+                label,
+                optional,
+            } = highlight;
+            let var = quote! { __miette_internal_var };
             if let Some(display) = label {
                 let (fmt, args) = display.expand_shorthand_cloned(&display_members);
+                if *optional {
+                    quote! {
+                        self.#span.clone().map(|#var|
+                            miette::LabeledSpan::new_with_span(
+                                std::option::Option::Some(format!(#fmt #args)),
+                                #var,
+                        ))
+                    }
+                } else {
+                    quote! {
+                        Some(miette::LabeledSpan::new_with_span(
+                            std::option::Option::Some(format!(#fmt #args)),
+                            self.#span.clone(),
+                        ))
+                    }
+                }
+            } else if *optional {
                 quote! {
-                    miette::LabeledSpan::new_with_span(
-                        std::option::Option::Some(format!(#fmt #args)),
-                        self.#span.clone(),
-                    )
+                    self.#span.clone().map(|#var|
+                        miette::LabeledSpan::new_with_span(
+                            std::option::Option::None,
+                            #var,
+                    ))
                 }
             } else {
                 quote! {
-                    miette::LabeledSpan::new_with_span(
+                    Some(miette::LabeledSpan::new_with_span(
                         std::option::Option::None,
                         self.#span.clone(),
-                    )
+                    ))
                 }
             }
         });
@@ -124,9 +172,9 @@ impl Labels {
             #[allow(unused_variables)]
             fn labels(&self) -> std::option::Option<std::boxed::Box<dyn std::iter::Iterator<Item = miette::LabeledSpan> + '_>> {
                 let Self #display_pat = self;
-                Some(Box::new(vec![
+                std::option::Option::Some(Box::new(vec![
                     #(#labels),*
-                ].into_iter()))
+                ].into_iter().filter(Option::is_some).map(Option::unwrap)))
             }
         })
     }
@@ -139,27 +187,46 @@ impl Labels {
                 let (display_pat, display_members) = display_pat_members(fields);
                 labels.as_ref().and_then(|labels| {
                 let variant_labels = labels.0.iter().map(|label| {
-                    let Label { span, label } = label;
+                    let Label { span, label, optional } = label;
                     let field = match &span {
                         syn::Member::Named(ident) => ident.clone(),
                         syn::Member::Unnamed(syn::Index { index, .. }) => {
                             format_ident!("_{}", index)
                         }
                     };
+                    let var = quote! { __miette_internal_var };
                     if let Some(display) = label {
                         let (fmt, args) = display.expand_shorthand_cloned(&display_members);
+                        if *optional {
+                            quote! {
+                                #field.clone().map(|#var|
+                                    miette::LabeledSpan::new_with_span(
+                                        std::option::Option::Some(format!(#fmt #args)),
+                                        #var,
+                                ))
+                            }
+                        } else {
+                            quote! {
+                                Some(miette::LabeledSpan::new_with_span(
+                                    std::option::Option::Some(format!(#fmt #args)),
+                                    #field.clone(),
+                                ))
+                            }
+                        }
+                    } else if *optional {
                         quote! {
-                            miette::LabeledSpan::new_with_span(
-                                std::option::Option::Some(format!(#fmt #args)),
-                                #field.clone(),
-                            )
+                            #field.clone().map(|#var|
+                                miette::LabeledSpan::new_with_span(
+                                    std::option::Option::None,
+                                    #var,
+                            ))
                         }
                     } else {
                         quote! {
-                            miette::LabeledSpan::new_with_span(
+                            Some(miette::LabeledSpan::new_with_span(
                                 std::option::Option::None,
                                 #field.clone(),
-                            )
+                            ))
                         }
                     }
                 });
@@ -169,7 +236,7 @@ impl Labels {
                     _ => Some(quote! {
                         Self::#variant_name #display_pat => std::option::Option::Some(std::boxed::Box::new(vec![
                             #(#variant_labels),*
-                        ].into_iter())),
+                        ].into_iter().filter(Option::is_some).map(Option::unwrap))),
                     }),
                 }
                 })
diff --git a/src/handlers/graphical.rs b/src/handlers/graphical.rs
index cd316cc..ecc297b 100644
--- a/src/handlers/graphical.rs
+++ b/src/handlers/graphical.rs
@@ -350,9 +350,7 @@ impl GraphicalReportHandler {
         let labels = labels
             .iter()
             .zip(self.theme.styles.highlights.iter().cloned().cycle())
-            .map(|(label, st)| {
-                FancySpan::new(label.label().map(String::from), label.inner().clone(), st)
-            })
+            .map(|(label, st)| FancySpan::new(label.label().map(String::from), *label.inner(), st))
             .collect::<Vec<_>>();
 
         // The max number of gutter-lines that will be active at any given
diff --git a/src/lib.rs b/src/lib.rs
index e40b7d8..af26647 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -368,6 +368,14 @@
 //!     // They'll be rendered sequentially.
 //!     #[label("This is bad")]
 //!     snip2: (usize, usize), // `(usize, usize)` is `Into<SourceSpan>`!
+//!
+//!     // Snippets can be optional, by tagging them as such:
+//!     #[label(optional, "some text")]
+//!     snip3: Option<SourceSpan>,
+//!
+//!     // with or without label text
+//!     #[label(optional)]
+//!     snip4: Option<SourceSpan>,
 //! }
 //! ```
 //!
diff --git a/tests/derive.rs b/tests/derive.rs
index e269621..1ee1bfa 100644
--- a/tests/derive.rs
+++ b/tests/derive.rs
@@ -309,6 +309,10 @@ fn test_snippet_named_struct() {
         var2: (usize, usize),
         #[label]
         var3: (usize, usize),
+        #[label(optional, "var 4")]
+        var4: Option<(usize, usize)>,
+        #[label(optional)]
+        var5: Option<(usize, usize)>,
     }
 }
 
@@ -322,6 +326,8 @@ fn test_snippet_unnamed_struct() {
         #[label("{0}")] SourceSpan,
         #[label = "idk"] SourceSpan,
         #[label] SourceSpan,
+        #[label(optional, "foo")] Option<SourceSpan>,
+        #[label(optional)] Option<SourceSpan>,
     );
 }
 
@@ -342,6 +348,10 @@ fn test_snippet_enum() {
             var1: SourceSpan,
             #[label]
             var2: SourceSpan,
+            #[label(optional, "var 3")]
+            var3: Option<(usize, usize)>,
+            #[label(optional)]
+            var4: Option<(usize, usize)>,
         },
         #[diagnostic(code(foo::b))]
         B(
@@ -350,6 +360,8 @@ fn test_snippet_enum() {
             #[label("{1}")] SourceSpan,
             #[label = "blorp"] SourceSpan,
             #[label] SourceSpan,
+            #[label(optional, "foo")] Option<SourceSpan>,
+            #[label(optional)] Option<SourceSpan>,
         ),
     }
 }