Repository: clap-rs/clap
Pull request #1875 (2020-04-28 09:40:39)

Implemented arg_enum for option and vec
<!--
If your PR closes some issues, please write `Closes #XXXX`
where `XXXX` is the number of the issue you want to fix.
Each issue goes on it's own line.
-->



Patch:
diff --git a/clap_derive/src/derives/from_argmatches.rs b/clap_derive/src/derives/from_argmatches.rs
index 31230f50..df7b2713 100644
--- a/clap_derive/src/derives/from_argmatches.rs
+++ b/clap_derive/src/derives/from_argmatches.rs
@@ -11,16 +11,17 @@
 // This work was derived from Structopt (https://github.com/TeXitoi/structopt)
 // commit#ea76fa1b1b273e65e3b0b1046643715b49bec51f which is licensed under the
 // MIT/Apache 2.0 license.
+use proc_macro2::TokenStream;
 use quote::{quote, quote_spanned};
-use syn::{punctuated::Punctuated, spanned::Spanned, Field, Ident, Token};
+use syn::{punctuated::Punctuated, spanned::Spanned, Field, Ident, Token, Type};
 
-use super::{sub_type, Attrs, Kind, ParserKind, Ty};
+use super::{sub_type, subty_if_name, Attrs, Kind, ParserKind, Ty};
 
 pub fn gen_for_struct(
     struct_name: &Ident,
     fields: &Punctuated<Field, Token![,]>,
     parent_attribute: &Attrs,
-) -> proc_macro2::TokenStream {
+) -> TokenStream {
     let constructor = gen_constructor(fields, parent_attribute);
 
     quote! {
@@ -44,7 +45,7 @@ pub fn gen_for_struct(
     }
 }
 
-pub fn gen_for_enum(name: &Ident) -> proc_macro2::TokenStream {
+pub fn gen_for_enum(name: &Ident) -> TokenStream {
     quote! {
         #[allow(dead_code, unreachable_code, unused_variables)]
         #[allow(
@@ -68,10 +69,18 @@ pub fn gen_for_enum(name: &Ident) -> proc_macro2::TokenStream {
     }
 }
 
+fn gen_arg_enum_parse(ty: &Type, attrs: &Attrs) -> TokenStream {
+    let ci = attrs.case_insensitive();
+
+    quote_spanned! { ty.span()=>
+        |s| <#ty as ::clap::ArgEnum>::from_str(s, #ci).unwrap()
+    }
+}
+
 pub fn gen_constructor(
     fields: &Punctuated<Field, Token![,]>,
     parent_attribute: &Attrs,
-) -> proc_macro2::TokenStream {
+) -> TokenStream {
     let fields = fields.iter().map(|field| {
         let attrs = Attrs::from_field(
             field,
@@ -117,7 +126,7 @@ pub fn gen_constructor(
                 let parser = attrs.parser();
                 let func = &parser.func;
                 let span = parser.kind.span();
-                let (value_of, values_of, parse) = match *parser.kind {
+                let (value_of, values_of, mut parse) = match *parser.kind {
                     FromStr => (
                         quote_spanned!(span=> value_of),
                         quote_spanned!(span=> values_of),
@@ -155,10 +164,18 @@ pub fn gen_constructor(
                         matches.is_present(#name)
                     },
 
-                    Ty::Option => quote_spanned! { ty.span()=>
-                        matches.#value_of(#name)
-                            .map(#parse)
-                    },
+                    Ty::Option => {
+                        if attrs.is_enum() {
+                            if let Some(subty) = subty_if_name(&field.ty, "Option") {
+                                parse = gen_arg_enum_parse(subty, &attrs);
+                            }
+                        }
+
+                        quote_spanned! { ty.span()=>
+                            matches.#value_of(#name)
+                                .map(#parse)
+                        }
+                    }
 
                     Ty::OptionOption => quote_spanned! { ty.span()=>
                         if matches.is_present(#name) {
@@ -178,11 +195,19 @@ pub fn gen_constructor(
                         }
                     },
 
-                    Ty::Vec => quote_spanned! { ty.span()=>
-                        matches.#values_of(#name)
-                            .map(|v| v.map(#parse).collect())
-                            .unwrap_or_else(Vec::new)
-                    },
+                    Ty::Vec => {
+                        if attrs.is_enum() {
+                            if let Some(subty) = subty_if_name(&field.ty, "Vec") {
+                                parse = gen_arg_enum_parse(subty, &attrs);
+                            }
+                        }
+
+                        quote_spanned! { ty.span()=>
+                            matches.#values_of(#name)
+                                .map(|v| v.map(#parse).collect())
+                                .unwrap_or_else(Vec::new)
+                        }
+                    }
 
                     Ty::Other if occurrences => quote_spanned! { ty.span()=>
                         #parse(matches.#value_of(#name))
@@ -193,16 +218,9 @@ pub fn gen_constructor(
                     },
 
                     Ty::Other => {
-                        let parse = if attrs.is_enum() {
-                            let field_ty = &field.ty;
-                            let ci = attrs.case_insensitive();
-
-                            quote_spanned! { field_ty.span()=>
-                                |s| <#field_ty as ::clap::ArgEnum>::from_str(s, #ci).unwrap()
-                            }
-                        } else {
-                            parse
-                        };
+                        if attrs.is_enum() {
+                            parse = gen_arg_enum_parse(&field.ty, &attrs);
+                        }
 
                         quote_spanned! { ty.span()=>
                             matches.#value_of(#name)
diff --git a/clap_derive/src/derives/into_app.rs b/clap_derive/src/derives/into_app.rs
index 85b34d36..b42d55b3 100644
--- a/clap_derive/src/derives/into_app.rs
+++ b/clap_derive/src/derives/into_app.rs
@@ -17,13 +17,12 @@ use std::env;
 use proc_macro2::TokenStream;
 use proc_macro_error::abort;
 use quote::{quote, quote_spanned};
-use syn::{punctuated::Punctuated, spanned::Spanned, Attribute, Field, Ident, Token};
+use syn::{punctuated::Punctuated, spanned::Spanned, Attribute, Field, Ident, Token, Type};
 
 use super::{
-    spanned::Sp, ty::Ty, Attrs, GenOutput, Kind, Name, ParserKind, DEFAULT_CASING,
-    DEFAULT_ENV_CASING,
+    spanned::Sp, sub_type, subty_if_name, ty::Ty, Attrs, GenOutput, Kind, Name, ParserKind,
+    DEFAULT_CASING, DEFAULT_ENV_CASING,
 };
-use crate::derives::ty::sub_type;
 
 pub fn gen_for_struct(
     struct_name: &Ident,
@@ -109,7 +108,7 @@ fn gen_into_app_fn(attrs: &[Attribute]) -> GenOutput {
 fn gen_augment_clap_fn(
     fields: &Punctuated<Field, Token![,]>,
     parent_attribute: &Attrs,
-) -> proc_macro2::TokenStream {
+) -> TokenStream {
     let app_var = Ident::new("app", proc_macro2::Span::call_site());
     let augmentation = gen_app_augmentation(fields, &app_var, parent_attribute);
     quote! {
@@ -119,13 +118,19 @@ fn gen_augment_clap_fn(
     }
 }
 
+fn gen_arg_enum_possible_values(ty: &Type) -> TokenStream {
+    quote_spanned! { ty.span()=>
+        .possible_values(&<#ty as ::clap::ArgEnum>::VARIANTS)
+    }
+}
+
 /// Generate a block of code to add arguments/subcommands corresponding to
 /// the `fields` to an app.
 pub fn gen_app_augmentation(
     fields: &Punctuated<Field, Token![,]>,
     app_var: &Ident,
     parent_attribute: &Attrs,
-) -> proc_macro2::TokenStream {
+) -> TokenStream {
     let mut subcmds = fields.iter().filter_map(|field| {
         let attrs = Attrs::from_field(
             &field,
@@ -195,6 +200,7 @@ pub fn gen_app_augmentation(
 
                 let parser = attrs.parser();
                 let func = &parser.func;
+
                 let validator = match *parser.kind {
                     _ if attrs.is_enum() => quote!(),
                     ParserKind::TryFromStr => quote_spanned! { func.span()=>
@@ -213,10 +219,21 @@ pub fn gen_app_augmentation(
                 let modifier = match **ty {
                     Ty::Bool => quote!(),
 
-                    Ty::Option => quote_spanned! { ty.span()=>
-                        .takes_value(true)
-                        #validator
-                    },
+                    Ty::Option => {
+                        let mut possible_values = quote!();
+
+                        if attrs.is_enum() {
+                            if let Some(subty) = subty_if_name(&field.ty, "Option") {
+                                possible_values = gen_arg_enum_possible_values(subty);
+                            }
+                        };
+
+                        quote_spanned! { ty.span()=>
+                            .takes_value(true)
+                            #possible_values
+                            #validator
+                        }
+                    }
 
                     Ty::OptionOption => quote_spanned! { ty.span()=>
                         .takes_value(true)
@@ -233,11 +250,22 @@ pub fn gen_app_augmentation(
                         #validator
                     },
 
-                    Ty::Vec => quote_spanned! { ty.span()=>
-                        .takes_value(true)
-                        .multiple(true)
-                        #validator
-                    },
+                    Ty::Vec => {
+                        let mut possible_values = quote!();
+
+                        if attrs.is_enum() {
+                            if let Some(subty) = subty_if_name(&field.ty, "Vec") {
+                                possible_values = gen_arg_enum_possible_values(subty);
+                            }
+                        };
+
+                        quote_spanned! { ty.span()=>
+                            .takes_value(true)
+                            .multiple(true)
+                            #possible_values
+                            #validator
+                        }
+                    }
 
                     Ty::Other if occurrences => quote_spanned! { ty.span()=>
                         .multiple_occurrences(true)
@@ -250,15 +278,10 @@ pub fn gen_app_augmentation(
 
                     Ty::Other => {
                         let required = !attrs.has_method("default_value");
+                        let mut possible_values = quote!();
 
-                        let possible_values = if attrs.is_enum() {
-                            let field_ty = &field.ty;
-
-                            quote_spanned! { field_ty.span()=>
-                                .possible_values(&<#field_ty as ::clap::ArgEnum>::VARIANTS)
-                            }
-                        } else {
-                            quote!()
+                        if attrs.is_enum() {
+                            possible_values = gen_arg_enum_possible_values(&field.ty);
                         };
 
                         quote_spanned! { ty.span()=>
diff --git a/clap_derive/src/derives/mod.rs b/clap_derive/src/derives/mod.rs
index aa891586..e3afdc33 100644
--- a/clap_derive/src/derives/mod.rs
+++ b/clap_derive/src/derives/mod.rs
@@ -23,10 +23,9 @@ pub mod spanned;
 mod subcommand;
 pub mod ty;
 
-// pub use self::arg_enum::derive_arg_enum;
 pub use self::attrs::{
     Attrs, CasingStyle, GenOutput, Kind, Name, Parser, ParserKind, DEFAULT_CASING,
     DEFAULT_ENV_CASING,
 };
 pub use self::clap::derive_clap;
-pub use self::ty::{sub_type, Ty};
+pub use self::ty::{sub_type, subty_if_name, Ty};
diff --git a/clap_derive/src/derives/ty.rs b/clap_derive/src/derives/ty.rs
index f9313f71..b046ac6d 100644
--- a/clap_derive/src/derives/ty.rs
+++ b/clap_derive/src/derives/ty.rs
@@ -80,7 +80,7 @@ where
         })
 }
 
-fn subty_if_name<'a>(ty: &'a syn::Type, name: &str) -> Option<&'a syn::Type> {
+pub fn subty_if_name<'a>(ty: &'a syn::Type, name: &str) -> Option<&'a syn::Type> {
     subty_if(ty, |seg| seg.ident == name)
 }
 


Test patch:
diff --git a/clap_derive/tests/arg_enum.rs b/clap_derive/tests/arg_enum.rs
index b61723e5..9b3ab65c 100644
--- a/clap_derive/tests/arg_enum.rs
+++ b/clap_derive/tests/arg_enum.rs
@@ -249,3 +249,63 @@ fn multiple_alias() {
         Opt::parse_from(&["", "t"])
     );
 }
+
+#[test]
+fn option() {
+    #[derive(Clap, PartialEq, Debug)]
+    enum ArgChoice {
+        Foo,
+        Bar,
+    }
+
+    #[derive(Clap, PartialEq, Debug)]
+    struct Opt {
+        #[clap(arg_enum)]
+        arg: Option<ArgChoice>,
+    };
+
+    assert_eq!(Opt { arg: None }, Opt::parse_from(&[""]));
+    assert_eq!(
+        Opt {
+            arg: Some(ArgChoice::Foo)
+        },
+        Opt::parse_from(&["", "foo"])
+    );
+    assert_eq!(
+        Opt {
+            arg: Some(ArgChoice::Bar)
+        },
+        Opt::parse_from(&["", "bar"])
+    );
+    assert!(Opt::try_parse_from(&["", "fOo"]).is_err());
+}
+
+#[test]
+fn vector() {
+    #[derive(Clap, PartialEq, Debug)]
+    enum ArgChoice {
+        Foo,
+        Bar,
+    }
+
+    #[derive(Clap, PartialEq, Debug)]
+    struct Opt {
+        #[clap(arg_enum, short, long)]
+        arg: Vec<ArgChoice>,
+    };
+
+    assert_eq!(Opt { arg: vec![] }, Opt::parse_from(&[""]));
+    assert_eq!(
+        Opt {
+            arg: vec![ArgChoice::Foo]
+        },
+        Opt::parse_from(&["", "-a", "foo"])
+    );
+    assert_eq!(
+        Opt {
+            arg: vec![ArgChoice::Foo, ArgChoice::Bar]
+        },
+        Opt::parse_from(&["", "-a", "foo", "bar"])
+    );
+    assert!(Opt::try_parse_from(&["", "-a", "fOo"]).is_err());
+}
diff --git a/clap_derive/tests/macro-errors.rs b/clap_derive/tests/ui.rs
similarity index 87%
rename from clap_derive/tests/macro-errors.rs
rename to clap_derive/tests/ui.rs
index 9c5bf437..54b405a0 100644
--- a/clap_derive/tests/macro-errors.rs
+++ b/clap_derive/tests/ui.rs
@@ -5,7 +5,7 @@
 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
 // option. This file may not be copied, modified, or distributed
 
-#[rustversion::attr(any(not(stable), before(1.42)), ignore)]
+#[rustversion::attr(any(not(stable), before(1.43)), ignore)]
 #[test]
 fn ui() {
     let t = trybuild::TestCases::new();