feat: Add 'ord' option for PyClass and corresponding tests

newsfragments/4202.added.md

@@ -0,0 +1 @@
+Support optional `ord` argument to `#[pyclass]` for simple enums in the form of `#[pyclass(ord)]`.

pyo3-macros-backend/src/attributes.rs

@@ -36,6 +36,7 @@ pub mod kw {
     syn::custom_keyword!(transparent);
     syn::custom_keyword!(unsendable);
     syn::custom_keyword!(weakref);
+    syn::custom_keyword!(ord);
 }
 
 #[derive(Clone, Debug)]

pyo3-macros-backend/src/pyclass.rs

@@ -72,6 +72,7 @@ pub struct PyClassPyO3Options {
     pub subclass: Option<kw::subclass>,
     pub unsendable: Option<kw::unsendable>,
     pub weakref: Option<kw::weakref>,
+    pub ord: Option<kw::ord>,
 }
 
 enum PyClassPyO3Option {
@@ -90,6 +91,7 @@ enum PyClassPyO3Option {
     Subclass(kw::subclass),
     Unsendable(kw::unsendable),
     Weakref(kw::weakref),
+    Ord(kw::ord),
 }
 
 impl Parse for PyClassPyO3Option {
@@ -125,6 +127,8 @@ impl Parse for PyClassPyO3Option {
             input.parse().map(PyClassPyO3Option::Unsendable)
         } else if lookahead.peek(attributes::kw::weakref) {
             input.parse().map(PyClassPyO3Option::Weakref)
+        } else if lookahead.peek(attributes::kw::ord) {
+            input.parse().map(PyClassPyO3Option::Ord)
         } else {
             Err(lookahead.error())
         }
@@ -179,6 +183,7 @@ impl PyClassPyO3Options {
             PyClassPyO3Option::Subclass(subclass) => set_option!(subclass),
             PyClassPyO3Option::Unsendable(unsendable) => set_option!(unsendable),
             PyClassPyO3Option::Weakref(weakref) => set_option!(weakref),
+            PyClassPyO3Option::Ord(ord) => set_option!(ord),
         }
         Ok(())
     }
@@ -775,6 +780,61 @@ fn impl_simple_enum(
         (int_impl, int_slot)
     };
 
+    // Add ordering comparisons if ord option is passed to pyclass
+    let richcmp_ord = match args.options.ord {
+        Some(_) => {
+            quote! {
+                #pyo3_path::basic::CompareOp::Gt => {
+                    let self_val = self.__pyo3__int__();
+                    if let Ok(i) = other.extract::<#repr_type>() {
+                        return Ok((self_val > i).to_object(py));
+                    }
+                    if let Ok(other) = other.extract::<#pyo3_path::PyRef<Self>>() {
+                        return Ok((self_val > other.__pyo3__int__()).to_object(py));
+                    }
+
+                    return Ok(py.NotImplemented());
+                }
+                #pyo3_path::basic::CompareOp::Lt => {
+                    let self_val = self.__pyo3__int__();
+                    if let Ok(i) = other.extract::<#repr_type>() {
+                        return Ok((self_val < i).to_object(py));
+                    }
+                    if let Ok(other) = other.extract::<#pyo3_path::PyRef<Self>>() {
+                        return Ok((self_val < other.__pyo3__int__()).to_object(py));
+                    }
+
+                    return Ok(py.NotImplemented());
+                }
+                #pyo3_path::basic::CompareOp::Le => {
+                    let self_val = self.__pyo3__int__();
+                    if let Ok(i) = other.extract::<#repr_type>() {
+                        return Ok((self_val <= i).to_object(py));
+                    }
+                    if let Ok(other) = other.extract::<#pyo3_path::PyRef<Self>>() {
+                        return Ok((self_val <= other.__pyo3__int__()).to_object(py));
+                    }
+
+                    return Ok(py.NotImplemented());
+                }
+                #pyo3_path::basic::CompareOp::Ge => {
+                    let self_val = self.__pyo3__int__();
+                    if let Ok(i) = other.extract::<#repr_type>() {
+                        return Ok((self_val >= i).to_object(py));
+                    }
+                    if let Ok(other) = other.extract::<#pyo3_path::PyRef<Self>>() {
+                        return Ok((self_val >= other.__pyo3__int__()).to_object(py));
+                    }
+
+                    return Ok(py.NotImplemented());
+                }
+            }
+        }
+        _ => {
+            quote! {}
+        }
+    };
+
     let (default_richcmp, default_richcmp_slot) = {
         let mut richcmp_impl: syn::ImplItemFn = syn::parse_quote! {
             fn __pyo3__richcmp__(
@@ -809,7 +869,10 @@ fn impl_simple_enum(
 
                         return Ok(py.NotImplemented());
                     }
-                    _ => Ok(py.NotImplemented()),
+                    #richcmp_ord
+                    _ => {
+                        Ok(py.NotImplemented())
+                    },
                 }
             }
         };

tests/test_enum.rs

@@ -13,6 +13,13 @@ pub enum MyEnum {
     OtherVariant,
 }
 
+#[pyclass(ord)]
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub enum MyEnumOrd {
+    Variant,
+    OtherVariant,
+}
+
 #[test]
 fn test_enum_class_attr() {
     Python::with_gil(|py| {
@@ -73,6 +80,30 @@ fn test_enum_eq_incomparable() {
     })
 }
 
+#[test]
+fn test_enum_ord_comparable_opt_in_only() {
+    Python::with_gil(|py| {
+        let var1 = Py::new(py, MyEnum::Variant).unwrap();
+        let var2 = Py::new(py, MyEnum::OtherVariant).unwrap();
+        // ordering on simple enums if opt in only, thus raising an error below
+        py_expect_exception!(py, var1 var2, "(var1 > var2) == False", PyTypeError);
+    })
+}
+
+#[test]
+fn test_enum_ord_comparable() {
+    Python::with_gil(|py| {
+        let var1 = Py::new(py, MyEnumOrd::Variant).unwrap();
+        let var2 = Py::new(py, MyEnumOrd::OtherVariant).unwrap();
+        let var3 = Py::new(py, MyEnumOrd::OtherVariant).unwrap();
+        py_assert!(py, var1 var2, "(var1 > var2) == False");
+        py_assert!(py, var1 var2, "(var1 < var2) == True");
+        py_assert!(py, var1 var2, "(var1 >= var2) == False");
+        py_assert!(py, var2 var3, "(var3 >= var2) == True");
+        py_assert!(py, var1 var2, "(var1 <= var2) == True");
+    })
+}
+
 #[pyclass]
 enum CustomDiscriminant {
     One = 1,

tests/ui/invalid_pyclass_args.stderr

@@ -1,4 +1,4 @@
-error: expected one of: `crate`, `dict`, `extends`, `freelist`, `frozen`, `get_all`, `mapping`, `module`, `name`, `rename_all`, `sequence`, `set_all`, `subclass`, `unsendable`, `weakref`
+error: expected one of: `crate`, `dict`, `extends`, `freelist`, `frozen`, `get_all`, `mapping`, `module`, `name`, `rename_all`, `sequence`, `set_all`, `subclass`, `unsendable`, `weakref`, `ord`
  --> tests/ui/invalid_pyclass_args.rs:3:11
   |
 3 | #[pyclass(extend=pyo3::types::PyDict)]
@@ -46,7 +46,7 @@ error: expected string literal
 24 | #[pyclass(module = my_module)]
    |                    ^^^^^^^^^
 
-error: expected one of: `crate`, `dict`, `extends`, `freelist`, `frozen`, `get_all`, `mapping`, `module`, `name`, `rename_all`, `sequence`, `set_all`, `subclass`, `unsendable`, `weakref`
+error: expected one of: `crate`, `dict`, `extends`, `freelist`, `frozen`, `get_all`, `mapping`, `module`, `name`, `rename_all`, `sequence`, `set_all`, `subclass`, `unsendable`, `weakref`, `ord`
   --> tests/ui/invalid_pyclass_args.rs:27:11
    |
 27 | #[pyclass(weakrev)]