From 88d00e0e20bf95e03b8f81dcd5adf38c917e190e Mon Sep 17 00:00:00 2001
From: Gavrilikhin Daniil <gavrilikhin.daniil@gmail.com>
Date: Sat, 21 Oct 2023 11:50:15 +0800
Subject: [PATCH] feat(derive): Allow optional sources in derive (#301)

Fixes: https://github.com/zkat/miette/issues/217
---
 miette-derive/src/source_code.rs | 42 +++++++++++++++++++++++++++++---
 tests/derive.rs                  | 42 ++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 3 deletions(-)

diff --git a/miette-derive/src/source_code.rs b/miette-derive/src/source_code.rs
index 62f28e76..e1b85ab4 100644
--- a/miette-derive/src/source_code.rs
+++ b/miette-derive/src/source_code.rs
@@ -10,6 +10,7 @@ use crate::{
 
 pub struct SourceCode {
     source_code: syn::Member,
+    is_option: bool,
 }
 
 impl SourceCode {
@@ -27,6 +28,19 @@ impl SourceCode {
         for (i, field) in fields.iter().enumerate() {
             for attr in &field.attrs {
                 if attr.path().is_ident("source_code") {
+                    let is_option = if let syn::Type::Path(syn::TypePath {
+                        path: syn::Path { segments, .. },
+                        ..
+                    }) = &field.ty
+                    {
+                        segments
+                            .last()
+                            .map(|seg| seg.ident == "Option")
+                            .unwrap_or(false)
+                    } else {
+                        false
+                    };
+
                     let source_code = if let Some(ident) = field.ident.clone() {
                         syn::Member::Named(ident)
                     } else {
@@ -35,7 +49,10 @@ impl SourceCode {
                             span: field.span(),
                         })
                     };
-                    return Ok(Some(SourceCode { source_code }));
+                    return Ok(Some(SourceCode {
+                        source_code,
+                        is_option,
+                    }));
                 }
             }
         }
@@ -45,11 +62,21 @@ impl SourceCode {
     pub(crate) fn gen_struct(&self, fields: &syn::Fields) -> Option<TokenStream> {
         let (display_pat, _display_members) = display_pat_members(fields);
         let src = &self.source_code;
+        let ret = if self.is_option {
+            quote! {
+                self.#src.as_ref().map(|s| s as _)
+            }
+        } else {
+            quote! {
+                Some(&self.#src)
+            }
+        };
+
         Some(quote! {
             #[allow(unused_variables)]
             fn source_code(&self) -> std::option::Option<&dyn miette::SourceCode> {
                 let Self #display_pat = self;
-                Some(&self.#src)
+                #ret
             }
         })
     }
@@ -68,10 +95,19 @@ impl SourceCode {
                         }
                     };
                     let variant_name = ident.clone();
+                    let ret = if source_code.is_option {
+                        quote! {
+                            #field.as_ref().map(|s| s as _)
+                        }
+                    } else {
+                        quote! {
+                            std::option::Option::Some(#field)
+                        }
+                    };
                     match &fields {
                         syn::Fields::Unit => None,
                         _ => Some(quote! {
-                            Self::#variant_name #display_pat => std::option::Option::Some(#field),
+                            Self::#variant_name #display_pat => #ret,
                         }),
                     }
                 })
diff --git a/tests/derive.rs b/tests/derive.rs
index f4739a7c..d1da1bb0 100644
--- a/tests/derive.rs
+++ b/tests/derive.rs
@@ -584,3 +584,45 @@ fn test_unit_enum_display() {
         "hello from unit help"
     );
 }
+
+#[test]
+fn test_optional_source_code() {
+    #[derive(Debug, Diagnostic, Error)]
+    #[error("struct with optional source")]
+    struct Struct {
+        #[source_code]
+        src: Option<String>,
+    }
+    assert!(Struct { src: None }.source_code().is_none());
+    assert!(Struct {
+        src: Some("".to_string())
+    }
+    .source_code()
+    .is_some());
+
+    #[derive(Debug, Diagnostic, Error)]
+    enum Enum {
+        #[error("variant1 with optional source")]
+        Variant1 {
+            #[source_code]
+            src: Option<String>,
+        },
+        #[error("variant2 with optional source")]
+        Variant2 {
+            #[source_code]
+            src: Option<String>,
+        },
+    }
+    assert!(Enum::Variant1 { src: None }.source_code().is_none());
+    assert!(Enum::Variant1 {
+        src: Some("".to_string())
+    }
+    .source_code()
+    .is_some());
+    assert!(Enum::Variant2 { src: None }.source_code().is_none());
+    assert!(Enum::Variant2 {
+        src: Some("".to_string())
+    }
+    .source_code()
+    .is_some());
+}