Merge pull request #2630 from epage/flatten

fix(derive): Allow partial update of flattened Subcommand arguments

clap_derive/src/derives/subcommand.rs

@@ -49,6 +49,7 @@ pub fn gen_for_enum(enum_name: &Ident, attrs: &[Attribute], e: &DataEnum) -> Tok
     );
     let augmentation = gen_augment(&e.variants, &attrs, false);
     let augmentation_update = gen_augment(&e.variants, &attrs, true);
+    let has_subcommand = gen_has_subcommand(&e.variants, &attrs);
 
     quote! {
         #from_arg_matches
@@ -72,6 +73,9 @@ pub fn gen_for_enum(enum_name: &Ident, attrs: &[Attribute], e: &DataEnum) -> Tok
             fn augment_subcommands_for_update <'b>(app: clap::App<'b>) -> clap::App<'b> {
                 #augmentation_update
             }
+            fn has_subcommand(name: &str) -> bool {
+                #has_subcommand
+            }
         }
     }
 }
@@ -135,8 +139,14 @@ fn gen_augment(
                 Kind::Flatten => match variant.fields {
                     Unnamed(FieldsUnnamed { ref unnamed, .. }) if unnamed.len() == 1 => {
                         let ty = &unnamed[0];
-                        quote! {
-                            let app = <#ty as clap::Subcommand>::augment_subcommands(app);
+                        if override_required {
+                            quote! {
+                                let app = <#ty as clap::Subcommand>::augment_subcommands_for_update(app);
+                            }
+                        } else {
+                            quote! {
+                                let app = <#ty as clap::Subcommand>::augment_subcommands(app);
+                            }
                         }
                     }
                     _ => abort!(
@@ -238,6 +248,74 @@ fn gen_augment(
     }
 }
 
+fn gen_has_subcommand(
+    variants: &Punctuated<Variant, Token![,]>,
+    parent_attribute: &Attrs,
+) -> TokenStream {
+    use syn::Fields::*;
+
+    let mut ext_subcmd = false;
+
+    let (flatten_variants, variants): (Vec<_>, Vec<_>) = variants
+        .iter()
+        .filter_map(|variant| {
+            let attrs = Attrs::from_variant(
+                variant,
+                parent_attribute.casing(),
+                parent_attribute.env_casing(),
+            );
+
+            if let Kind::ExternalSubcommand = &*attrs.kind() {
+                ext_subcmd = true;
+                None
+            } else {
+                Some((variant, attrs))
+            }
+        })
+        .partition(|(_, attrs)| {
+            let kind = attrs.kind();
+            matches!(&*kind, Kind::Flatten)
+        });
+
+    let match_arms = variants.iter().map(|(_variant, attrs)| {
+        let sub_name = attrs.cased_name();
+        quote! {
+            #sub_name => true,
+        }
+    });
+    let child_subcommands = flatten_variants
+        .iter()
+        .map(|(variant, _attrs)| match variant.fields {
+            Unnamed(ref fields) if fields.unnamed.len() == 1 => {
+                let ty = &fields.unnamed[0];
+                quote! {
+                    if <#ty as clap::Subcommand>::has_subcommand(name) {
+                        return true;
+                    }
+                }
+            }
+            _ => abort!(
+                variant,
+                "`flatten` is usable only with single-typed tuple variants"
+            ),
+        });
+
+    if ext_subcmd {
+        quote! { true }
+    } else {
+        quote! {
+            match name {
+                #( #match_arms )*
+                _ => {
+                    #( #child_subcommands )else*
+
+                    false
+                }
+            }
+        }
+    }
+}
+
 fn gen_from_arg_matches(
     name: &Ident,
     variants: &Punctuated<Variant, Token![,]>,
@@ -390,7 +468,7 @@ fn gen_update_from_arg_matches(
 ) -> TokenStream {
     use syn::Fields::*;
 
-    let variants: Vec<_> = variants
+    let (flatten, variants): (Vec<_>, Vec<_>) = variants
         .iter()
         .filter_map(|variant| {
             let attrs = Attrs::from_variant(
@@ -401,11 +479,14 @@ fn gen_update_from_arg_matches(
 
             match &*attrs.kind() {
                 // Fallback to `from_arg_matches`
-                Kind::ExternalSubcommand | Kind::Flatten => None,
+                Kind::ExternalSubcommand => None,
                 _ => Some((variant, attrs)),
             }
         })
-        .collect();
+        .partition(|(_, attrs)| {
+            let kind = attrs.kind();
+            matches!(&*kind, Kind::Flatten)
+        });
 
     let subcommands = variants.iter().map(|(variant, attrs)| {
         let sub_name = attrs.cased_name();
@@ -454,6 +535,27 @@ fn gen_update_from_arg_matches(
         }
     });
 
+    let child_subcommands = flatten.iter().map(|(variant, _attrs)| {
+        let variant_name = &variant.ident;
+        match variant.fields {
+            Unnamed(ref fields) if fields.unnamed.len() == 1 => {
+                let ty = &fields.unnamed[0];
+                quote! {
+                    if <#ty as clap::Subcommand>::has_subcommand(name) {
+                        if let #name :: #variant_name (child) = s {
+                            <#ty as clap::FromArgMatches>::update_from_arg_matches(child, arg_matches);
+                            return;
+                        }
+                    }
+                }
+            }
+            _ => abort!(
+                variant,
+                "`flatten` is usable only with single-typed tuple variants"
+            ),
+        }
+    });
+
     quote! {
         fn update_from_arg_matches<'b>(
             &mut self,
@@ -463,6 +565,7 @@ fn gen_update_from_arg_matches(
                 match (name, self) {
                     #( #subcommands ),*
                     (other_name, s) => {
+                        #( #child_subcommands )*
                         if let Some(sub) = <Self as clap::FromArgMatches>::from_arg_matches(arg_matches) {
                             *s = sub;
                         }

clap_derive/tests/flatten.rs

@@ -107,6 +107,7 @@ enum BaseCli {
 #[derive(Clap, PartialEq, Debug)]
 struct Command1 {
     arg1: i32,
+    arg2: i32,
 }
 
 #[derive(Clap, PartialEq, Debug)]
@@ -124,15 +125,31 @@ enum Opt {
 #[test]
 fn merge_subcommands_with_flatten() {
     assert_eq!(
-        Opt::BaseCli(BaseCli::Command1(Command1 { arg1: 42 })),
-        Opt::parse_from(&["test", "command1", "42"])
+        Opt::BaseCli(BaseCli::Command1(Command1 { arg1: 42, arg2: 44 })),
+        Opt::parse_from(&["test", "command1", "42", "44"])
     );
     assert_eq!(
         Opt::Command2(Command2 { arg2: 43 }),
         Opt::parse_from(&["test", "command2", "43"])
     );
 }
 
+#[test]
+fn update_subcommands_with_flatten() {
+    let mut opt = Opt::BaseCli(BaseCli::Command1(Command1 { arg1: 12, arg2: 14 }));
+    opt.try_update_from(&["test", "command1", "42", "44"])
+        .unwrap();
+    assert_eq!(Opt::parse_from(&["test", "command1", "42", "44"]), opt);
+
+    let mut opt = Opt::BaseCli(BaseCli::Command1(Command1 { arg1: 12, arg2: 14 }));
+    opt.try_update_from(&["test", "command1", "42"]).unwrap();
+    assert_eq!(Opt::parse_from(&["test", "command1", "42", "14"]), opt);
+
+    let mut opt = Opt::BaseCli(BaseCli::Command1(Command1 { arg1: 12, arg2: 14 }));
+    opt.try_update_from(&["test", "command2", "43"]).unwrap();
+    assert_eq!(Opt::parse_from(&["test", "command2", "43"]), opt);
+}
+
 #[test]
 fn flatten_with_doc_comment() {
     #[derive(Clap, Debug)]
@@ -151,12 +168,3 @@ fn flatten_with_doc_comment() {
         opts: DaemonOpts,
     }
 }
-
-#[test]
-fn update_subcommands_with_flatten() {
-    let mut opt = Opt::BaseCli(BaseCli::Command1(Command1 { arg1: 12 }));
-    opt.update_from(&["test", "command1", "42"]);
-    assert_eq!(Opt::parse_from(&["test", "command1", "42"]), opt);
-    opt.update_from(&["test", "command2", "43"]);
-    assert_eq!(Opt::parse_from(&["test", "command2", "43"]), opt);
-}

src/derive.rs

@@ -254,6 +254,8 @@ pub trait Subcommand: FromArgMatches + Sized {
     ///
     /// See also [`IntoApp`].
     fn augment_subcommands_for_update(app: App<'_>) -> App<'_>;
+    /// Test whether `Self` can parse a specific subcommand
+    fn has_subcommand(name: &str) -> bool;
 }
 
 /// Parse arguments into enums.
@@ -354,4 +356,7 @@ impl<T: Subcommand> Subcommand for Box<T> {
     fn augment_subcommands_for_update(app: App<'_>) -> App<'_> {
         <T as Subcommand>::augment_subcommands_for_update(app)
     }
+    fn has_subcommand(name: &str) -> bool {
+        <T as Subcommand>::has_subcommand(name)
+    }
 }