Implement __int__ and __eq__ betweeen enum and int

Author: jovenlin0527

pyo3-macros-backend/src/pyclass.rs

@@ -501,30 +501,58 @@ fn impl_enum_class(
         }
     };
 
+    let repr_type = &enum_.repr;
+
+    let default_int = {
+        // 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, }
+        });
+        quote! {
+            #[doc(hidden)]
+            #[allow(non_snake_case)]
+            #[pyo3(name = "__int__")]
+            fn __pyo3__int__(&self) -> #repr_type {
+                match self {
+                    #(#variants_to_int)*
+                }
+            }
+        }
+    };
+
     let default_richcmp = {
         let variants_eq = variants.iter().map(|variant| {
             let variant_name = variant.ident;
-            quote! {(#cls::#variant_name, #cls::#variant_name) => true.to_object(py),}
+            quote! {(#cls::#variant_name, #cls::#variant_name) => Ok(true.to_object(py)),}
         });
         quote! {
             #[doc(hidden)]
             #[allow(non_snake_case)]
             #[pyo3(name = "__richcmp__")]
-            fn __pyo3__richcmp__(&self, py: ::pyo3::Python, other: &Self, op: ::pyo3::basic::CompareOp) -> PyObject {
+            fn __pyo3__richcmp__(&self, py: ::pyo3::Python, other: &PyAny, op: ::pyo3::basic::CompareOp) -> PyResult<PyObject> {
                 match op {
                     ::pyo3::basic::CompareOp::Eq => {
+                        if let Ok(i) = other.extract::<#repr_type>() {
+                            let self_val = self.__pyo3__int__();
+                            return Ok((self_val == i).to_object(py));
+                        }
+                        let other = other.extract::<PyRef<Self>>()?;
+                        let other = &*other;
                         match (self, other) {
                             #(#variants_eq)*
-                            _ => false.to_object(py),
+                            _ => Ok(false.to_object(py)),
                         }
                     }
-                    _ => py.NotImplemented(),
+                    _ => Ok(py.NotImplemented()),
                 }
             }
         }
     };
 
-    let default_impls = gen_default_slot_impls(cls, vec![default_repr_impl, default_richcmp]);
+    let default_impls =
+        gen_default_slot_impls(cls, vec![default_repr_impl, default_richcmp, default_int]);
+
     Ok(quote! {
 
         #pytypeinfo

tests/test_enum.rs

@@ -90,3 +90,59 @@ fn test_custom_discriminant() {
         "#);
     })
 }
+
+#[test]
+fn test_enum_to_int() {
+    Python::with_gil(|py| {
+        let one = Py::new(py, CustomDiscriminant::One).unwrap();
+        py_assert!(py, one, "int(one) == 1");
+        let test_noexception = Py::new(py, MyEnum::Variant).unwrap();
+        py_run!(py, test_noexception, "int(test_noexception)");
+    })
+}
+
+#[test]
+fn test_enum_compare_int() {
+    Python::with_gil(|py| {
+        let one = Py::new(py, CustomDiscriminant::One).unwrap();
+        py_run!(
+            py,
+            one,
+            r#"
+            assert one == 1
+            assert 1 == one
+            assert one != 2
+        "#
+        )
+    })
+}
+
+#[pyclass]
+#[repr(u8)]
+enum SmallEnum {
+    V = 1,
+}
+
+#[test]
+fn test_enum_compare_int_no_throw_when_overflow() {
+    Python::with_gil(|py| {
+        let v = Py::new(py, SmallEnum::V).unwrap();
+        py_assert!(py, v, "v != 1<<30")
+    })
+}
+
+#[pyclass]
+#[repr(usize)]
+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");
+    })
+}