cfg features for enum variants

newsfragments/4509.fixed.md

@@ -0,0 +1 @@
+Fix compile failure when using `#[cfg]` attributes for simple enum variants.

pyo3-macros-backend/src/pyclass.rs

@@ -17,7 +17,7 @@ use crate::attributes::{
 use crate::konst::{ConstAttributes, ConstSpec};
 use crate::method::{FnArg, FnSpec, PyArg, RegularArg};
 use crate::pyfunction::ConstructorAttribute;
-use crate::pyimpl::{gen_py_const, PyClassMethodsType};
+use crate::pyimpl::{gen_py_const, get_cfg_attributes, PyClassMethodsType};
 use crate::pymethod::{
     impl_py_getter_def, impl_py_setter_def, MethodAndMethodDef, MethodAndSlotDef, PropertyType,
     SlotDef, __GETITEM__, __HASH__, __INT__, __LEN__, __REPR__, __RICHCMP__, __STR__,
@@ -533,7 +533,12 @@ impl<'a> PyClassSimpleEnum<'a> {
                 _ => bail_spanned!(variant.span() => "Must be a unit variant."),
             };
             let options = EnumVariantPyO3Options::take_pyo3_options(&mut variant.attrs)?;
-            Ok(PyClassEnumUnitVariant { ident, options })
+            let cfg_attrs = get_cfg_attributes(&variant.attrs);
+            Ok(PyClassEnumUnitVariant {
+                ident,
+                options,
+                cfg_attrs,
+            })
         }
 
         let ident = &enum_.ident;
@@ -693,6 +698,7 @@ impl<'a> EnumVariant for PyClassEnumVariant<'a> {
 struct PyClassEnumUnitVariant<'a> {
     ident: &'a syn::Ident,
     options: EnumVariantPyO3Options,
+    cfg_attrs: Vec<&'a syn::Attribute>,
 }
 
 impl<'a> EnumVariant for PyClassEnumUnitVariant<'a> {
@@ -877,20 +883,23 @@ fn impl_simple_enum(
         ensure_spanned!(variant.options.constructor.is_none(), variant.options.constructor.span() => "`constructor` can't be used on a simple enum variant");
     }
 
+    let variant_cfg_check = generate_cfg_check(&variants, cls);
+
     let (default_repr, default_repr_slot) = {
         let variants_repr = variants.iter().map(|variant| {
             let variant_name = variant.ident;
+            let cfg_attrs = &variant.cfg_attrs;
             // Assuming all variants are unit variants because they are the only type we support.
             let repr = format!(
                 "{}.{}",
                 get_class_python_name(cls, args),
                 variant.get_python_name(args),
             );
-            quote! { #cls::#variant_name => #repr, }
+            quote! { #(#cfg_attrs)* #cls::#variant_name => #repr, }
         });
         let mut repr_impl: syn::ImplItemFn = syn::parse_quote! {
             fn __pyo3__repr__(&self) -> &'static str {
-                match self {
+                match *self {
                     #(#variants_repr)*
                 }
             }
@@ -908,11 +917,12 @@ fn impl_simple_enum(
         // This implementation allows us to convert &T to #repr_type without implementing `Copy`
         let variants_to_int = variants.iter().map(|variant| {
             let variant_name = variant.ident;
-            quote! { #cls::#variant_name => #cls::#variant_name as #repr_type, }
+            let cfg_attrs = &variant.cfg_attrs;
+            quote! { #(#cfg_attrs)* #cls::#variant_name => #cls::#variant_name as #repr_type, }
         });
         let mut int_impl: syn::ImplItemFn = syn::parse_quote! {
             fn __pyo3__int__(&self) -> #repr_type {
-                match self {
+                match *self {
                     #(#variants_to_int)*
                 }
             }
@@ -936,7 +946,9 @@ fn impl_simple_enum(
         methods_type,
         simple_enum_default_methods(
             cls,
-            variants.iter().map(|v| (v.ident, v.get_python_name(args))),
+            variants
+                .iter()
+                .map(|v| (v.ident, v.get_python_name(args), &v.cfg_attrs)),
             ctx,
         ),
         default_slots,
@@ -945,6 +957,8 @@ fn impl_simple_enum(
     .impl_all(ctx)?;
 
     Ok(quote! {
+        #variant_cfg_check
+
         #pytypeinfo
 
         #pyclass_impls
@@ -1474,7 +1488,13 @@ fn generate_default_protocol_slot(
 
 fn simple_enum_default_methods<'a>(
     cls: &'a syn::Ident,
-    unit_variant_names: impl IntoIterator<Item = (&'a syn::Ident, Cow<'a, syn::Ident>)>,
+    unit_variant_names: impl IntoIterator<
+        Item = (
+            &'a syn::Ident,
+            Cow<'a, syn::Ident>,
+            &'a Vec<&'a syn::Attribute>,
+        ),
+    >,
     ctx: &Ctx,
 ) -> Vec<MethodAndMethodDef> {
     let cls_type = syn::parse_quote!(#cls);
@@ -1490,7 +1510,25 @@ fn simple_enum_default_methods<'a>(
     };
     unit_variant_names
         .into_iter()
-        .map(|(var, py_name)| gen_py_const(&cls_type, &variant_to_attribute(var, &py_name), ctx))
+        .map(|(var, py_name, attrs)| {
+            let method = gen_py_const(&cls_type, &variant_to_attribute(var, &py_name), ctx);
+            let associated_method_tokens = method.associated_method;
+            let method_def_tokens = method.method_def;
+
+            let associated_method = quote! {
+                #(#attrs)*
+                #associated_method_tokens
+            };
+            let method_def = quote! {
+                #(#attrs)*
+                #method_def_tokens
+            };
+
+            MethodAndMethodDef {
+                associated_method,
+                method_def,
+            }
+        })
         .collect()
 }
 
@@ -2395,6 +2433,34 @@ fn define_inventory_class(inventory_class_name: &syn::Ident, ctx: &Ctx) -> Token
     }
 }
 
+fn generate_cfg_check(variants: &[PyClassEnumUnitVariant<'_>], cls: &syn::Ident) -> TokenStream {
+    if variants.is_empty() {
+        return quote! {};
+    }
+
+    let mut conditions = Vec::new();
+
+    for variant in variants {
+        let cfg_attrs = &variant.cfg_attrs;
+
+        if cfg_attrs.is_empty() {
+            return quote! {};
+        }
+
+        for attr in cfg_attrs {
+            if let syn::Meta::List(meta) = &attr.meta {
+                let cfg_tokens = &meta.tokens;
+                conditions.push(quote! { not(#cfg_tokens) });
+            }
+        }
+    }
+
+    quote! {
+        #[cfg(all(#(#conditions),*))]
+        ::core::compile_error!(concat!("All variants of enum `", stringify!(#cls), "` have been disabled by cfg attributes"));
+    }
+}
+
 const UNIQUE_GET: &str = "`get` may only be specified once";
 const UNIQUE_SET: &str = "`set` may only be specified once";
 const UNIQUE_NAME: &str = "`name` may only be specified once";

pyo3-macros-backend/src/pyimpl.rs

@@ -330,7 +330,7 @@ fn submit_methods_inventory(
     }
 }
 
-fn get_cfg_attributes(attrs: &[syn::Attribute]) -> Vec<&syn::Attribute> {
+pub(crate) fn get_cfg_attributes(attrs: &[syn::Attribute]) -> Vec<&syn::Attribute> {
     attrs
         .iter()
         .filter(|attr| attr.path().is_ident("cfg"))

tests/test_field_cfg.rs

@@ -17,6 +17,15 @@ struct CfgClass {
     pub b: u32,
 }
 
+#[pyclass(eq, eq_int)]
+#[derive(PartialEq)]
+enum CfgSimpleEnum {
+    #[cfg(any())]
+    DisabledVariant,
+    #[cfg(not(any()))]
+    EnabledVariant,
+}
+
 #[test]
 fn test_cfg() {
     Python::with_gil(|py| {
@@ -27,3 +36,18 @@ fn test_cfg() {
         assert_eq!(b, 3);
     });
 }
+
+#[test]
+fn test_cfg_simple_enum() {
+    Python::with_gil(|py| {
+        let simple = py.get_type::<CfgSimpleEnum>();
+        pyo3::py_run!(
+            py,
+            simple,
+            r#"
+            assert hasattr(simple, "EnabledVariant")
+            assert not hasattr(simple, "DisabledVariant")
+        "#
+        );
+    })
+}

tests/ui/invalid_pyclass_enum.rs

@@ -93,4 +93,13 @@ enum InvalidOrderedComplexEnum2 {
     VariantB { msg: String }
 }
 
+#[pyclass(eq)]
+#[derive(PartialEq)]
+enum AllEnumVariantsDisabled {
+    #[cfg(any())]
+    DisabledA,
+    #[cfg(not(all()))]
+    DisabledB,
+}
+
 fn main() {}

tests/ui/invalid_pyclass_enum.stderr

@@ -66,6 +66,14 @@ error: The `ord` option requires the `eq` option.
 83 | #[pyclass(ord)]
    |           ^^^
 
+error: All variants of enum `AllEnumVariantsDisabled` have been disabled by cfg attributes
+  --> tests/ui/invalid_pyclass_enum.rs:96:1
+   |
+96 | #[pyclass(eq)]
+   | ^^^^^^^^^^^^^^
+   |
+   = note: this error originates in the attribute macro `pyclass` (in Nightly builds, run with -Z macro-backtrace for more info)
+
 error[E0369]: binary operation `==` cannot be applied to type `&SimpleEqOptRequiresPartialEq`
   --> tests/ui/invalid_pyclass_enum.rs:31:11
    |