From 6a70c744e02ee8ab0afb7a1449b7b3a5d4aac98c Mon Sep 17 00:00:00 2001 From: Ed Page Date: Thu, 15 Jul 2021 12:36:14 -0500 Subject: [PATCH] fix(derive): Allow partial update of flattened Subcommand arguments When using `#[clap(flatten)]` inside of a `Subcommand`, we would do a `from` instead of an `update`. The challenge is knowing when we are going into a flattened subcommand vs changing the variant. To resolve this, I added a `Subcommand:has_subcommand(name)` trait method that we generate, so we can ask. --- clap_derive/src/derives/subcommand.rs | 113 ++++++++++++++++++++++++-- clap_derive/tests/flatten.rs | 30 ++++--- src/derive.rs | 5 ++ 3 files changed, 132 insertions(+), 16 deletions(-) diff --git a/clap_derive/src/derives/subcommand.rs b/clap_derive/src/derives/subcommand.rs index 4b5bdd5d7d8..af0f7854620 100644 --- a/clap_derive/src/derives/subcommand.rs +++ b/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, + 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, @@ -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) = ::from_arg_matches(arg_matches) { *s = sub; } diff --git a/clap_derive/tests/flatten.rs b/clap_derive/tests/flatten.rs index 27419041702..febca659835 100644 --- a/clap_derive/tests/flatten.rs +++ b/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,8 +125,8 @@ 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 }), @@ -133,6 +134,22 @@ fn merge_subcommands_with_flatten() { ); } +#[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); -} diff --git a/src/derive.rs b/src/derive.rs index 4320e7d21aa..cc4dc52b6fc 100644 --- a/src/derive.rs +++ b/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 Subcommand for Box { fn augment_subcommands_for_update(app: App<'_>) -> App<'_> { ::augment_subcommands_for_update(app) } + fn has_subcommand(name: &str) -> bool { + ::has_subcommand(name) + } }