Implement choice enums (#196)

* Add derive macro for FromArgValue for choice enums

* Add basic documentation on choice enums

* Fix and add choice name override tests

* Update github workflow for new example

Author: kovaxis

.github/workflows/rust.yml

@@ -15,7 +15,7 @@ jobs:
       run: cargo test --workspace --verbose
 
     - name: Run simple_example
-      run: cargo run --package argh --example simple_example two --fooey
+      run: cargo run --package argh --example simple_example two --fooey --woot quiet
 
   rustfmt:
     name: rustfmt

argh/examples/simple_example.rs

@@ -34,6 +34,16 @@ struct SubCommandTwo {
     #[argh(switch)]
     /// whether to fooey
     fooey: bool,
+
+    #[argh(option)]
+    /// how to woot
+    woot: Woot,
+}
+
+#[derive(argh::FromArgValue, PartialEq, Debug)]
+enum Woot {
+    Quiet,
+    Loud,
 }
 
 fn main() {

argh/src/lib.rs

@@ -111,6 +111,30 @@
 //! }
 //! ```
 //!
+//! `FromArgValue` can be automatically derived for `enum`s, with automatic
+//! error messages:
+//!
+//! ```
+//! use argh::{FromArgs, FromArgValue};
+//!
+//! #[derive(FromArgValue)]
+//! enum Mode {
+//!     SoftCore,
+//!     HardCore,
+//! }
+//!
+//! #[derive(FromArgs)]
+//! /// Do the thing.
+//! struct DoIt {
+//!     #[argh(option)]
+//!     /// how to do it
+//!     how: Mode,
+//! }
+//!
+//! // ./some_bin --how whatever
+//! // > Error parsing option '--how' with value 'whatever': expected "soft_core" or "hard_core"
+//! ```
+//!
 //! Positional arguments can be declared using `#[argh(positional)]`.
 //! These arguments will be parsed in order of their declaration in
 //! the structure:
@@ -319,7 +343,7 @@
 
 use std::str::FromStr;
 
-pub use argh_derive::{ArgsInfo, FromArgs};
+pub use argh_derive::{ArgsInfo, FromArgValue, FromArgs};
 
 /// Information about a particular command used for output.
 pub type CommandInfo = argh_shared::CommandInfo<'static>;
@@ -751,7 +775,7 @@ pub fn cargo_from_env<T: TopLevelCommand>() -> T {
 /// Any field type declared in a struct that derives `FromArgs` must implement
 /// this trait. A blanket implementation exists for types implementing
 /// `FromStr<Error: Display>`. Custom types can implement this trait
-/// directly.
+/// directly. It can also be derived on plain `enum`s without associated data.
 pub trait FromArgValue: Sized {
     /// Construct the type from a commandline value, returning an error string
     /// on failure.

argh/tests/lib.rs

@@ -11,7 +11,10 @@
     clippy::unwrap_in_result
 )]
 
-use {argh::FromArgs, std::fmt::Debug};
+use {
+    argh::{FromArgValue, FromArgs},
+    std::fmt::Debug,
+};
 
 #[test]
 fn basic_example() {
@@ -519,6 +522,87 @@ Woot
 Options:
   --option-name     fooey
   --help, help      display usage information
+"###,
+        );
+    }
+
+    /// Test choices
+    #[derive(FromArgs, PartialEq, Debug)]
+    struct WithChoices {
+        /// first choice with a default
+        #[argh(option, default = "TwoChoices::Chao")]
+        choice1: TwoChoices,
+        /// second choice.
+        #[argh(option)]
+        choice2: ThreeChoices,
+    }
+
+    #[derive(FromArgValue, PartialEq, Debug)]
+    enum TwoChoices {
+        Hola,
+        Chao,
+    }
+
+    #[derive(FromArgValue, PartialEq, Debug)]
+    enum ThreeChoices {
+        FirstChoice,
+        #[argh(name = "に")]
+        Two,
+        Three,
+    }
+
+    #[test]
+    fn with_choices() {
+        assert_output(
+            &["--choice2", "three"],
+            WithChoices { choice1: TwoChoices::Chao, choice2: ThreeChoices::Three },
+        );
+    }
+
+    #[test]
+    fn with_choices_snake_case() {
+        assert_output(
+            &["--choice2", "first_choice"],
+            WithChoices { choice1: TwoChoices::Chao, choice2: ThreeChoices::FirstChoice },
+        )
+    }
+
+    #[test]
+    fn override_default() {
+        assert_output(
+            &["--choice2", "first_choice", "--choice1", "hola"],
+            WithChoices { choice1: TwoChoices::Hola, choice2: ThreeChoices::FirstChoice },
+        )
+    }
+
+    #[test]
+    fn with_name_override() {
+        assert_output(
+            &["--choice2", "に", "--choice1", "hola"],
+            WithChoices { choice1: TwoChoices::Hola, choice2: ThreeChoices::Two },
+        )
+    }
+
+    #[test]
+    fn invalid_choice() {
+        assert_error::<WithChoices>(
+            &["--choice2", "something"],
+            r###"Error parsing option '--choice2' with value 'something': expected "first_choice", "に" or "three"
+"###,
+        )
+    }
+
+    #[test]
+    fn choice_help() {
+        assert_help_string::<WithChoices>(
+            r###"Usage: test_arg_0 [--choice1 <choice1>] --choice2 <choice2>
+
+Test choices
+
+Options:
+  --choice1         first choice with a default
+  --choice2         second choice.
+  --help, help      display usage information
 "###,
         );
     }

argh_derive/src/lib.rs

@@ -3,6 +3,7 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
+use parse_attrs::has_argh_attrs;
 use syn::ext::IdentExt as _;
 
 /// Implementation of the `FromArgs` and `argh(...)` derive attributes.
@@ -34,6 +35,14 @@ pub fn argh_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     gen.into()
 }
 
+/// Entrypoint for `#[derive(FromArgValue)]`.
+#[proc_macro_derive(FromArgValue, attributes(argh))]
+pub fn argh_value_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    let ast = syn::parse_macro_input!(input as syn::DeriveInput);
+    let gen = impl_from_arg_value(&ast);
+    gen.into()
+}
+
 /// Entrypoint for `#[derive(ArgsInfo)]`.
 #[proc_macro_derive(ArgsInfo, attributes(argh))]
 pub fn args_info_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
@@ -63,6 +72,25 @@ fn impl_from_args(input: &syn::DeriveInput) -> TokenStream {
     output_tokens
 }
 
+fn impl_from_arg_value(input: &syn::DeriveInput) -> TokenStream {
+    let errors = &Errors::default();
+    let mut output_tokens = match &input.data {
+        syn::Data::Enum(de) => impl_from_arg_value_enum(errors, &input.ident, &input.generics, de),
+        _ => {
+            errors.err(input, "`#[derive(FromArgValue)]` can only be applied to `enum`s");
+            TokenStream::new()
+        }
+    };
+    if has_argh_attrs(&input.attrs) {
+        errors.err(
+            &input.ident,
+            "`#[derive(FromArgValue)]` `enum`s do not support `#[argh(...)]` attributes",
+        );
+    }
+    errors.to_tokens(&mut output_tokens);
+    output_tokens
+}
+
 /// The kind of optionality a parameter has.
 enum Optionality {
     None,
@@ -1178,3 +1206,101 @@ fn enum_only_single_field_unnamed_variants<'a>(
         }
     }
 }
+
+/// Implements `FromArgValue` for a `#![derive(FromArgValue)]` enum (a choice enum).
+fn impl_from_arg_value_enum(
+    errors: &Errors,
+    name: &syn::Ident,
+    generic_args: &syn::Generics,
+    de: &syn::DataEnum,
+) -> TokenStream {
+    // An enum variant like `<name>`
+    struct ChoiceVariant<'a> {
+        ident: &'a syn::Ident,
+        name: syn::LitStr,
+    }
+
+    let variants: Vec<ChoiceVariant<'_>> = de
+        .variants
+        .iter()
+        .map(|variant| {
+            let ident = &variant.ident;
+            choice_enum_only_fieldless_variant(errors, &variant.fields);
+            let attrs = parse_attrs::ChoiceVariantAttrs::parse(errors, variant);
+            let name = match attrs.name_override {
+                Some(lit) => lit,
+                None => {
+                    let name_str = pascal_to_snake_case(&format!("{}", ident));
+                    syn::LitStr::new(&name_str, ident.span())
+                }
+            };
+            ChoiceVariant { ident, name }
+        })
+        .collect();
+
+    if variants.is_empty() {
+        errors.err(&de.variants, "Choice enums must have at least one variant");
+    }
+
+    let name_repeating = std::iter::repeat(name.clone());
+    let variant_idents = variants.iter().map(|x| x.ident);
+    let variant_names = variants.iter().map(|x| &x.name).collect::<Vec<_>>();
+    let err_literal = {
+        let mut err = "expected ".to_string();
+        for (i, name) in variant_names.iter().enumerate() {
+            if i == 0 {
+            } else if i == variant_names.len() - 1 {
+                err.push_str(" or ");
+            } else {
+                err.push_str(", ");
+            }
+            err.push_str(&format!("{:?}", name.value()));
+        }
+        LitStr::new(&err, name.span())
+    };
+    let (impl_generics, ty_generics, where_clause) = generic_args.split_for_impl();
+    quote! {
+        impl #impl_generics argh::FromArgValue for #name #ty_generics #where_clause {
+            fn from_arg_value(value: &str)
+                -> std::result::Result<Self, String>
+            {
+                Ok(match value {
+                    #(
+                        #variant_names => #name_repeating::#variant_idents,
+                    )*
+                    _ => {
+                        return Err(#err_literal.to_owned())
+                    }
+                })
+            }
+        }
+    }
+}
+
+/// Generates an error if the variant is not a field-less variant like `Foo`.
+fn choice_enum_only_fieldless_variant(errors: &Errors, variant_fields: &syn::Fields) {
+    match variant_fields {
+        syn::Fields::Unit => {}
+        _ => {
+            errors.err(
+                variant_fields,
+                "Choice `enum`s tagged with `#![derive(FromArgValue)]` do not support variants with associated data.",
+            );
+        }
+    }
+}
+
+fn pascal_to_snake_case(camel: &str) -> String {
+    let mut out = String::with_capacity(camel.len() + 8);
+    for (i, c) in camel.chars().enumerate() {
+        if c.is_uppercase() {
+            if i != 0 {
+                out.push('_');
+            }
+            out.extend(c.to_lowercase());
+        } else {
+            out.push(c);
+        }
+    }
+    out
+}