feat(derive): Support 'Option' when flattening

clap-rs · Oct 5, 2022 · 06d2049 · 06d2049
1 parent 33db640
commit 06d2049
Show file tree

Hide file tree

Showing 2 changed files with 92 additions and 11 deletions.
diff --git a/clap_derive/src/derives/args.rs b/clap_derive/src/derives/args.rs
@@ -202,23 +202,27 @@ pub fn gen_augment(
                         #implicit_methods;
                 })
             }
-            Kind::Flatten(_) => {
-                let ty = &field.ty;
+            Kind::Flatten(ty) => {
+                let inner_type = match (**ty, sub_type(&field.ty)) {
+                    (Ty::Option, Some(sub_type)) => sub_type,
+                    _ => &field.ty,
+                };
+
                 let next_help_heading = item.next_help_heading();
                 let next_display_order = item.next_display_order();
                 if override_required {
                     Some(quote_spanned! { kind.span()=>
                         let #app_var = #app_var
                             #next_help_heading
                             #next_display_order;
-                        let #app_var = <#ty as clap::Args>::augment_args_for_update(#app_var);
+                        let #app_var = <#inner_type as clap::Args>::augment_args_for_update(#app_var);
                     })
                 } else {
                     Some(quote_spanned! { kind.span()=>
                         let #app_var = #app_var
                             #next_help_heading
                             #next_display_order;
-                        let #app_var = <#ty as clap::Args>::augment_args(#app_var);
+                        let #app_var = <#inner_type as clap::Args>::augment_args(#app_var);
                     })
                 }
             }
@@ -430,14 +434,30 @@ pub fn gen_constructor(fields: &[(&Field, Item)]) -> TokenStream {
             }
 
             Kind::Flatten(ty) => {
+                let inner_type = match (**ty, sub_type(&field.ty)) {
+                    (Ty::Option, Some(sub_type)) => sub_type,
+                    _ => &field.ty,
+                };
                 match **ty {
                     Ty::Other => {
                         quote_spanned! { kind.span()=>
-                            #field_name: clap::FromArgMatches::from_arg_matches_mut(#arg_matches)?
+                            #field_name: <#inner_type as clap::FromArgMatches>::from_arg_matches_mut(#arg_matches)?
+                        }
+                    },
+                    Ty::Option => {
+                        quote_spanned! { kind.span()=>
+                            #field_name: {
+                                if <#inner_type as clap::Args>::group_id().map(|group_id| #arg_matches.contains_id(group_id.as_str())).unwrap_or(false) {
+                                    Some(
+                                        <#inner_type as clap::FromArgMatches>::from_arg_matches_mut(#arg_matches)?
+                                    )
+                                } else {
+                                    None
+                                }
+                            }
                         }
                     },
                     Ty::Vec |
-                    Ty::Option |
                     Ty::OptionOption |
                     Ty::OptionVec => {
                         abort!(
@@ -522,11 +542,29 @@ pub fn gen_updater(fields: &[(&Field, Item)], use_self: bool) -> TokenStream {
                 }
             }
 
-            Kind::Flatten(_) => {
-                let updater = quote_spanned! { kind.span()=> {
-                        #access
-                        clap::FromArgMatches::update_from_arg_matches_mut(#field_name, #arg_matches)?;
-                    }
+            Kind::Flatten(ty) => {
+                let inner_type = match (**ty, sub_type(&field.ty)) {
+                    (Ty::Option, Some(sub_type)) => sub_type,
+                    _ => &field.ty,
+                };
+
+                let updater = quote_spanned! { ty.span()=>
+                    <#inner_type as clap::FromArgMatches>::update_from_arg_matches_mut(#field_name, #arg_matches)?;
+                };
+
+                let updater = match **ty {
+                    Ty::Option => quote_spanned! { kind.span()=>
+                        if let Some(#field_name) = #field_name.as_mut() {
+                            #updater
+                        } else {
+                            *#field_name = Some(<#inner_type as clap::FromArgMatches>::from_arg_matches_mut(
+                                #arg_matches
+                            )?);
+                        }
+                    },
+                    _ => quote_spanned! { kind.span()=>
+                        #updater
+                    },
                 };
 
                 quote_spanned! { kind.span()=>

diff --git a/tests/derive/flatten.rs b/tests/derive/flatten.rs
@@ -255,3 +255,46 @@ fn docstrings_ordering_with_multiple_clap_partial() {
 
     assert!(short_help.contains("This is the docstring for Flattened"));
 }
+
+#[test]
+fn optional_flatten() {
+    #[derive(Parser, Debug, PartialEq, Eq)]
+    struct Opt {
+        #[command(flatten)]
+        source: Option<Source>,
+    }
+
+    #[derive(clap::Args, Debug, PartialEq, Eq)]
+    struct Source {
+        crates: Vec<String>,
+        #[arg(long)]
+        path: Option<std::path::PathBuf>,
+        #[arg(long)]
+        git: Option<String>,
+    }
+
+    assert_eq!(
+        Opt { source: None },
+        Opt::try_parse_from(&["test"]).unwrap()
+    );
+    assert_eq!(
+        Opt {
+            source: Some(Source {
+                crates: vec!["serde".to_owned()],
+                path: None,
+                git: None,
+            }),
+        },
+        Opt::try_parse_from(&["test", "serde"]).unwrap()
+    );
+    assert_eq!(
+        Opt {
+            source: Some(Source {
+                crates: Vec::new(),
+                path: Some("./".into()),
+                git: None,
+            }),
+        },
+        Opt::try_parse_from(&["test", "--path=./"]).unwrap()
+    );
+}