Parse #[repr(..)] for #[pyclass] enums. Allow custom discriminants.

Author: jovenlin0527

pyo3-macros-backend/src/pyclass.rs

@@ -406,6 +406,55 @@ struct PyClassEnumVariant<'a> {
     /* currently have no more options */
 }
 
+struct PyClassEnum<'a> {
+    ident: &'a syn::Ident,
+    // The underlying #[repr] of the enum, used to implement __int__ and __richcmp__.
+    // This matters when the underlying representation may not fit in `isize`.
+    #[allow(unused, dead_code)]
+    repr_type: syn::Ident,
+    variants: Vec<PyClassEnumVariant<'a>>,
+}
+
+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)
+        }
+        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_type = syn::Ident::new("isize", proc_macro2::Span::call_site());
+        if let Some(attr) = enum_.attrs.iter().find(|attr| attr.path.is_ident("repr")) {
+            let args =
+                attr.parse_args_with(Punctuated::<TokenStream, Token![!]>::parse_terminated)?;
+            if let Some(ident) = args
+                .into_iter()
+                .filter_map(|ts| syn::parse2::<syn::Ident>(ts).ok())
+                .find(is_numeric_type)
+            {
+                repr_type = ident;
+            }
+        }
+
+        let variants = enum_
+            .variants
+            .iter()
+            .map(extract_variant_data)
+            .collect::<syn::Result<_>>()?;
+        Ok(Self {
+            ident,
+            repr_type,
+            variants,
+        })
+    }
+}
+
 pub fn build_py_enum(
     enum_: &mut syn::ItemEnum,
     args: &PyClassArgs,
@@ -416,41 +465,37 @@ 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, options)
-}
-
-fn impl_enum(
-    enum_: &syn::ItemEnum,
-    args: &PyClassArgs,
-    variants: Vec<PyClassEnumVariant>,
-    methods_type: PyClassMethodsType,
-    options: PyClassPyO3Options,
-) -> syn::Result<TokenStream> {
-    let enum_name = &enum_.ident;
     let doc = utils::get_doc(
         &enum_.attrs,
         options
             .text_signature
             .as_ref()
             .map(|attr| (get_class_python_name(&enum_.ident, args), attr)),
     );
+    let enum_ = PyClassEnum::new(enum_)?;
+    impl_enum(enum_, args, doc, method_type, options)
+}
+
+fn impl_enum(
+    enum_: PyClassEnum,
+    args: &PyClassArgs,
+    doc: PythonDoc,
+    methods_type: PyClassMethodsType,
+    options: PyClassPyO3Options,
+) -> syn::Result<TokenStream> {
     let krate = get_pyo3_crate(&options.krate);
-    impl_enum_class(enum_name, args, variants, doc, methods_type, krate)
+    impl_enum_class(enum_, args, doc, methods_type, krate)
 }
 
 fn impl_enum_class(
-    cls: &syn::Ident,
+    enum_: PyClassEnum,
     args: &PyClassArgs,
-    variants: Vec<PyClassEnumVariant>,
     doc: PythonDoc,
     methods_type: PyClassMethodsType,
     krate: syn::Path,
 ) -> syn::Result<TokenStream> {
+    let cls = enum_.ident;
+    let variants = enum_.variants;
     let pytypeinfo = impl_pytypeinfo(cls, args, None);
     let pyclass_impls = PyClassImplsBuilder::new(cls, args, methods_type)
         .doc(doc)
@@ -528,9 +573,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

@@ -63,3 +63,52 @@ 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
+        "#);
+    })
+}
+
+#[pyclass]
+#[repr(usize)]
+#[allow(clippy::enum_clike_unportable_variant)]
+enum BigEnum {
+    V = usize::MAX,
+}
+
+#[test]
+fn test_big_enum_no_overflow() {
+    Python::with_gil(|py| {
+        let usize_max = usize::MAX;
+        let v = Py::new(py, BigEnum::V).unwrap();
+        py_assert!(py, usize_max v, "v == usize_max");
+        py_assert!(py, usize_max v, "int(v) == usize_max");
+    })
+}
+
+#[pyclass]
+#[repr(u16, align(8))]
+enum TestReprParse {
+    V,
+}
+
+#[test]
+fn test_repr_parse() {
+    assert_eq!(std::mem::align_of::<TestReprParse>(), 8);
+}