Repository: clap-rs/clap
Pull request #3591 (2022-03-30 01:24:10)

Fix ArgEnum non-unit variant
Hi!

As promised, I tackled the issue #3543, which I reported.
Please tell me what you think, I will refactor it if needed.


Patch:
diff --git a/clap_derive/src/derives/arg_enum.rs b/clap_derive/src/derives/arg_enum.rs
index 8b020cb4..77067e90 100644
--- a/clap_derive/src/derives/arg_enum.rs
+++ b/clap_derive/src/derives/arg_enum.rs
@@ -15,11 +15,11 @@ use crate::{
 };
 
 use proc_macro2::{Span, TokenStream};
-use proc_macro_error::abort_call_site;
+use proc_macro_error::{abort, abort_call_site};
 use quote::quote;
 use syn::{
-    punctuated::Punctuated, token::Comma, Attribute, Data, DataEnum, DeriveInput, Fields, Ident,
-    Variant,
+    punctuated::Punctuated, spanned::Spanned, token::Comma, Attribute, Data, DataEnum, DeriveInput,
+    Fields, Ident, Variant,
 };
 
 pub fn derive_arg_enum(input: &DeriveInput) -> TokenStream {
@@ -34,10 +34,6 @@ pub fn derive_arg_enum(input: &DeriveInput) -> TokenStream {
 }
 
 pub fn gen_for_enum(name: &Ident, attrs: &[Attribute], e: &DataEnum) -> TokenStream {
-    if !e.variants.iter().all(|v| matches!(v.fields, Fields::Unit)) {
-        return quote!();
-    };
-
     let attrs = Attrs::from_struct(
         Span::call_site(),
         attrs,
@@ -86,6 +82,9 @@ fn lits(
             if let Kind::Skip(_) = &*attrs.kind() {
                 None
             } else {
+                if !matches!(variant.fields, Fields::Unit) {
+                    abort!(variant.span(), "`#[derive(ArgEnum)]` only supports non-unit variants, unless they are skipped");
+                }
                 let fields = attrs.field_methods(false);
                 let name = attrs.cased_name();
                 Some((
diff --git a/examples/derive_ref/README.md b/examples/derive_ref/README.md
index 2b0955fe..e9d4ac19 100644
--- a/examples/derive_ref/README.md
+++ b/examples/derive_ref/README.md
@@ -82,6 +82,7 @@ fn main() {
 - `Subcommand` defines available subcommands.
   - Subcommand arguments can be defined in a struct-variant or automatically flattened with a tuple-variant.
 - `ArgEnum` allows parsing a value directly into an `enum`, erroring on unsupported values.
+  - The derive doesn't work on enums that contain non-unit variants, unless they are skipped
 
 See also the [tutorial](../tutorial_derive/README.md) and [examples](../README.md).
 


Test patch:
diff --git a/tests/derive/arg_enum.rs b/tests/derive/arg_enum.rs
index 957e870d..1cd8ee0d 100644
--- a/tests/derive/arg_enum.rs
+++ b/tests/derive/arg_enum.rs
@@ -321,6 +321,37 @@ fn skip_variant() {
     }
 }
 
+#[test]
+fn skip_non_unit_variant() {
+    #[derive(clap::ArgEnum, PartialEq, Debug, Clone)]
+    #[allow(dead_code)] // silence warning about `Baz` being unused
+    enum ArgChoice {
+        Foo,
+        Bar,
+        #[clap(skip)]
+        Baz(usize),
+    }
+
+    assert_eq!(
+        <ArgChoice as clap::ArgEnum>::value_variants()
+            .iter()
+            .map(clap::ArgEnum::to_possible_value)
+            .map(Option::unwrap)
+            .collect::<Vec<_>>(),
+        vec![
+            clap::PossibleValue::new("foo"),
+            clap::PossibleValue::new("bar")
+        ]
+    );
+
+    {
+        use clap::ArgEnum;
+        assert!(ArgChoice::from_str("foo", true).is_ok());
+        assert!(ArgChoice::from_str("bar", true).is_ok());
+        assert!(ArgChoice::from_str("baz", true).is_err());
+    }
+}
+
 #[test]
 fn from_str_invalid() {
     #[derive(clap::ArgEnum, PartialEq, Debug, Clone)]
diff --git a/tests/derive_ui/arg_enum_non_unit.rs b/tests/derive_ui/arg_enum_non_unit.rs
new file mode 100644
index 00000000..be031005
--- /dev/null
+++ b/tests/derive_ui/arg_enum_non_unit.rs
@@ -0,0 +1,10 @@
+use clap::ArgEnum;
+
+#[derive(ArgEnum, Clone, Debug)]
+enum Opt {
+    Foo(usize),
+}
+
+fn main() {
+    println!("{:?}", Opt::Foo(42));
+}
diff --git a/tests/derive_ui/arg_enum_non_unit.stderr b/tests/derive_ui/arg_enum_non_unit.stderr
new file mode 100644
index 00000000..117cfcec
--- /dev/null
+++ b/tests/derive_ui/arg_enum_non_unit.stderr
@@ -0,0 +1,5 @@
+error: `#[derive(ArgEnum)]` only supports non-unit variants, unless they are skipped
+ --> tests/derive_ui/arg_enum_non_unit.rs:5:5
+  |
+5 |     Foo(usize),
+  |     ^^^