Parse #[repr(..)] for #[pyclass] enums.

Author: jovenlin0527

pyo3-macros-backend/src/pyclass.rs

@@ -381,6 +381,66 @@ struct PyClassEnumVariant<'a> {
     /* currently have no more options */
 }
 
+struct PyClassEnum<'a> {
+    ident: &'a syn::Ident,
+    // The underyling representation of the enum.
+    // It's used to implement __int__ and __richcmp__.
+    // This matters when the underyling representation may not fit in `isize`.
+    #[allow(unused, dead_code)]
+    repr: syn::Ident,
+    variants: Vec<PyClassEnumVariant<'a>>,
+    doc: PythonDoc,
+}
+
+impl<'a> PyClassEnum<'a> {
+    fn new(enum_: &'a syn::ItemEnum) -> syn::Result<Self> {
+        fn is_numeric_type(t: &syn::Ident) -> bool {
+            [
+                "u8", "i8", "u16", "i16", "u32", "i32", "u64", "i64", "u128", "i128", "usize",
+                "isize",
+            ]
+            .iter()
+            .any(|&s| t == s)
+        }
+        struct Reprs(syn::punctuated::Punctuated<syn::Ident, Token![,]>);
+        impl Parse for Reprs {
+            fn parse(input: ParseStream) -> Result<Self> {
+                let inner = Punctuated::parse_terminated(input)?;
+                Ok(Self(inner))
+            }
+        }
+        let ident = &enum_.ident;
+        // According to the [reference](https://doc.rust-lang.org/reference/items/enumerations.html),
+        // "Under the default representation, the specified discriminant is interpreted as an isize
+        // value", so `isize` should be enough by default.
+        let mut repr = syn::Ident::new("isize", proc_macro2::Span::call_site());
+        for attr in &enum_.attrs {
+            if attr.path.is_ident("repr") {
+                let reprs: Reprs = attr.parse_args()?;
+                for r in reprs.0 {
+                    if is_numeric_type(&r) {
+                        repr = r;
+                        break;
+                    }
+                }
+            }
+        }
+        let doc = utils::get_doc(&enum_.attrs, None);
+
+        let variants = enum_
+            .variants
+            .iter()
+            .map(extract_variant_data)
+            .collect::<syn::Result<_>>()?;
+        Ok(Self {
+            ident,
+            repr,
+            variants,
+            doc,
+        })
+    }
+}
+
 pub fn build_py_enum(
     enum_: &syn::ItemEnum,
     args: PyClassArgs,
@@ -389,38 +449,32 @@ pub fn build_py_enum(
     if enum_.variants.is_empty() {
         bail_spanned!(enum_.brace_token.span => "Empty enums can't be #[pyclass].");
     }
-    let variants: Vec<PyClassEnumVariant> = enum_
-        .variants
-        .iter()
-        .map(extract_variant_data)
-        .collect::<syn::Result<_>>()?;
-    impl_enum(enum_, args, variants, method_type)
+    let enum_ = PyClassEnum::new(enum_)?;
+    impl_enum(enum_, args, method_type)
 }
 
 fn impl_enum(
-    enum_: &syn::ItemEnum,
-    attrs: PyClassArgs,
-    variants: Vec<PyClassEnumVariant>,
+    enum_: PyClassEnum,
+    args: PyClassArgs,
     methods_type: PyClassMethodsType,
 ) -> syn::Result<TokenStream> {
-    let enum_name = &enum_.ident;
-    let doc = utils::get_doc(&enum_.attrs, None);
-    let enum_cls = impl_enum_class(enum_name, &attrs, variants, doc, methods_type)?;
+    let enum_cls = impl_enum_class(enum_, &args, methods_type)?;
 
     Ok(quote! {
         #enum_cls
     })
 }
 
 fn impl_enum_class(
-    cls: &syn::Ident,
-    attr: &PyClassArgs,
-    variants: Vec<PyClassEnumVariant>,
-    doc: PythonDoc,
+    enum_: PyClassEnum,
+    args: &PyClassArgs,
     methods_type: PyClassMethodsType,
 ) -> syn::Result<TokenStream> {
-    let pytypeinfo = impl_pytypeinfo(cls, attr, None);
-    let pyclass_impls = PyClassImplsBuilder::new(cls, attr, methods_type)
+    let cls = enum_.ident;
+    let doc = enum_.doc;
+    let variants = enum_.variants;
+    let pytypeinfo = impl_pytypeinfo(cls, args, None);
+    let pyclass_impls = PyClassImplsBuilder::new(cls, args, methods_type)
         .doc(doc)
         .impl_all();
     let descriptors = unit_variants_as_descriptors(cls, variants.iter().map(|v| v.ident));
@@ -494,9 +548,6 @@ fn extract_variant_data(variant: &syn::Variant) -> syn::Result<PyClassEnumVarian
         Fields::Unit => &variant.ident,
         _ => bail_spanned!(variant.span() => "Currently only support unit variants."),
     };
-    if let Some(discriminant) = variant.discriminant.as_ref() {
-        bail_spanned!(discriminant.0.span() => "Currently does not support discriminats.")
-    };
     Ok(PyClassEnumVariant { ident })
 }
 

tests/test_enum.rs

@@ -61,3 +61,24 @@ fn test_default_repr_correct() {
         py_assert!(py, var2, "repr(var2) == 'MyEnum.OtherVariant'");
     })
 }
+
+#[pyclass]
+enum CustomDiscriminant {
+    One = 1,
+    Two = 2,
+}
+
+#[test]
+fn test_custom_discriminant() {
+    Python::with_gil(|py| {
+        #[allow(non_snake_case)]
+        let CustomDiscriminant = py.get_type::<CustomDiscriminant>();
+        let one = Py::new(py, CustomDiscriminant::One).unwrap();
+        let two = Py::new(py, CustomDiscriminant::Two).unwrap();
+        py_run!(py, CustomDiscriminant one two, r#"
+        assert CustomDiscriminant.One == one
+        assert CustomDiscriminant.Two == two
+        assert one != two
+        "#);
+    })
+}