feat: Add 'ord' option for PyClass and corresponding tests (#4202)

* feat: Add 'ord' option for PyClass and corresponding tests Updated the macros back-end to include 'ord' as an option for PyClass allowing for Python-style ordering comparison of enum variants. Additionally, test cases to verify the proper functioning of this new feature have been introduced. * update: fix formatting with cargo fmt * update: documented added feature in newsfragments * update: updated saved errors for comparison test for invalid pyclass args * update: removed nested match arms and extended cases for ordering instead * update: alphabetically ordered entries * update: added section to class documentation with example for using ord argument. * refactor: reduced duplication of code using closure to process tokens. * update: used ensure_spanned macro to emit compile time errors for uses of ord on complex enums or structs, updated test errors for bad compile cases * fix: remove errant character * update: added note about PartialOrd being required. * feat: implemented ordering for structs and complex enums. Retained the equality logic for simple enums until PartialEq is deprecated. * update: adjusted compile time error checks for missing PartialOrd implementations. Refactored growing set of comparison tests for simple and complex enums and structs into separate test file. * fix: updated with clippy findings * update: added not to pyclass parameters on ord (assumes that eq will be implemented and merged first) * update: rebased on main after merging of `eq` feature * update: format update * update: update all test output and doc tests * Update guide/src/class.md Co-authored-by: Icxolu <[email protected]> * Update pyo3-macros-backend/src/pyclass.rs Co-authored-by: Icxolu <[email protected]> * Update newsfragments/4202.added.md Co-authored-by: Icxolu <[email protected]> * Update guide/pyclass-parameters.md Co-authored-by: Icxolu <[email protected]> * update: added note about `ord` implementation with example. * fix doc formatting --------- Co-authored-by: Michael Gilbert <[email protected]> Co-authored-by: Icxolu <[email protected]>
PyO3 · Jun 7, 2024 · b8fb367 · b8fb367
1 parent fbb6f20
commit b8fb367
Show file tree

Hide file tree

Showing 12 changed files with 465 additions and 32 deletions.
diff --git a/guide/pyclass-parameters.md b/guide/pyclass-parameters.md
@@ -15,6 +15,7 @@
 | `mapping` |  Inform PyO3 that this class is a [`Mapping`][params-mapping], and so leave its implementation of sequence C-API slots empty. |
 | <span style="white-space: pre">`module = "module_name"`</span> |  Python code will see the class as being defined in this module. Defaults to `builtins`. |
 | <span style="white-space: pre">`name = "python_name"`</span> | Sets the name that Python sees this class as. Defaults to the name of the Rust struct. |
+| `ord` | Implements `__lt__`, `__gt__`, `__le__`, & `__ge__` using the `PartialOrd` implementation of the underlying Rust datatype. *Requires `eq`* |
 | `rename_all = "renaming_rule"` | Applies renaming rules to every getters and setters of a struct, or every variants of an enum. Possible values are: "camelCase", "kebab-case", "lowercase", "PascalCase", "SCREAMING-KEBAB-CASE", "SCREAMING_SNAKE_CASE", "snake_case", "UPPERCASE". |
 | `sequence` |  Inform PyO3 that this class is a [`Sequence`][params-sequence], and so leave its C-API mapping length slot empty. |
 | `set_all` | Generates setters for all fields of the pyclass. |

diff --git a/guide/src/class.md b/guide/src/class.md
@@ -1160,6 +1160,32 @@ Python::with_gil(|py| {
 })
 ```
 
+Ordering of enum variants is optionally added using `#[pyo3(ord)]`.  
+*Note: Implementation of the `PartialOrd` trait is required when passing the `ord` argument.  If not implemented, a compile time error is raised.*
+
+```rust
+# use pyo3::prelude::*;
+#[pyclass(eq, ord)]
+#[derive(PartialEq, PartialOrd)]
+enum MyEnum{
+    A,
+    B,
+    C,
+}
+
+Python::with_gil(|py| {
+    let cls = py.get_type_bound::<MyEnum>();
+    let a = Py::new(py, MyEnum::A).unwrap();
+    let b = Py::new(py, MyEnum::B).unwrap();
+    let c = Py::new(py, MyEnum::C).unwrap();
+    pyo3::py_run!(py, cls a b c, r#"
+        assert (a < b) == True
+        assert (c <= b) == False
+        assert (c > a) == True
+    "#)
+})
+```
+
 You may not use enums as a base class or let enums inherit from other classes.
 
 ```rust,compile_fail

diff --git a/guide/src/class/object.md b/guide/src/class/object.md
@@ -249,6 +249,16 @@ To implement `__eq__` using the Rust [`PartialEq`] trait implementation, the `eq
 struct Number(i32);
 ```
 
+To implement `__lt__`, `__le__`, `__gt__`, & `__ge__` using the Rust `PartialOrd` trait implementation, the `ord` option can be used. *Note: Requires `eq`.*
+
+```rust
+# use pyo3::prelude::*;
+#
+#[pyclass(eq, ord)]
+#[derive(PartialEq, PartialOrd)]
+struct Number(i32);
+```
+
 ### Truthyness
 
 We'll consider `Number` to be `True` if it is nonzero:

diff --git a/newsfragments/4202.added.md b/newsfragments/4202.added.md
@@ -0,0 +1 @@
+Add `#[pyclass(ord)]` to implement ordering based on `PartialOrd`.
diff --git a/pyo3-macros-backend/src/attributes.rs b/pyo3-macros-backend/src/attributes.rs
@@ -29,6 +29,7 @@ pub mod kw {
     syn::custom_keyword!(mapping);
     syn::custom_keyword!(module);
     syn::custom_keyword!(name);
+    syn::custom_keyword!(ord);
     syn::custom_keyword!(pass_module);
     syn::custom_keyword!(rename_all);
     syn::custom_keyword!(sequence);

diff --git a/pyo3-macros-backend/src/pyclass.rs b/pyo3-macros-backend/src/pyclass.rs
@@ -70,6 +70,7 @@ pub struct PyClassPyO3Options {
     pub mapping: Option<kw::mapping>,
     pub module: Option<ModuleAttribute>,
     pub name: Option<NameAttribute>,
+    pub ord: Option<kw::ord>,
     pub rename_all: Option<RenameAllAttribute>,
     pub sequence: Option<kw::sequence>,
     pub set_all: Option<kw::set_all>,
@@ -91,6 +92,7 @@ pub enum PyClassPyO3Option {
     Mapping(kw::mapping),
     Module(ModuleAttribute),
     Name(NameAttribute),
+    Ord(kw::ord),
     RenameAll(RenameAllAttribute),
     Sequence(kw::sequence),
     SetAll(kw::set_all),
@@ -126,6 +128,8 @@ impl Parse for PyClassPyO3Option {
             input.parse().map(PyClassPyO3Option::Module)
         } else if lookahead.peek(kw::name) {
             input.parse().map(PyClassPyO3Option::Name)
+        } else if lookahead.peek(attributes::kw::ord) {
+            input.parse().map(PyClassPyO3Option::Ord)
         } else if lookahead.peek(kw::rename_all) {
             input.parse().map(PyClassPyO3Option::RenameAll)
         } else if lookahead.peek(attributes::kw::sequence) {
@@ -189,6 +193,7 @@ impl PyClassPyO3Options {
             PyClassPyO3Option::Mapping(mapping) => set_option!(mapping),
             PyClassPyO3Option::Module(module) => set_option!(module),
             PyClassPyO3Option::Name(name) => set_option!(name),
+            PyClassPyO3Option::Ord(ord) => set_option!(ord),
             PyClassPyO3Option::RenameAll(rename_all) => set_option!(rename_all),
             PyClassPyO3Option::Sequence(sequence) => set_option!(sequence),
             PyClassPyO3Option::SetAll(set_all) => set_option!(set_all),
@@ -1665,7 +1670,10 @@ fn impl_pytypeinfo(
     }
 }
 
-fn pyclass_richcmp_arms(options: &PyClassPyO3Options, ctx: &Ctx) -> TokenStream {
+fn pyclass_richcmp_arms(
+    options: &PyClassPyO3Options,
+    ctx: &Ctx,
+) -> std::result::Result<TokenStream, syn::Error> {
     let Ctx { pyo3_path } = ctx;
 
     let eq_arms = options
@@ -1684,9 +1692,34 @@ fn pyclass_richcmp_arms(options: &PyClassPyO3Options, ctx: &Ctx) -> TokenStream
         })
         .unwrap_or_default();
 
-    // TODO: `ord` can be integrated here (#4202)
-    #[allow(clippy::let_and_return)]
-    eq_arms
+    if let Some(ord) = options.ord {
+        ensure_spanned!(options.eq.is_some(), ord.span() => "The `ord` option requires the `eq` option.");
+    }
+
+    let ord_arms = options
+        .ord
+        .map(|ord| {
+            quote_spanned! { ord.span() =>
+                #pyo3_path::pyclass::CompareOp::Gt => {
+                    ::std::result::Result::Ok(#pyo3_path::conversion::IntoPy::into_py(self_val > other, py))
+                },
+                #pyo3_path::pyclass::CompareOp::Lt => {
+                    ::std::result::Result::Ok(#pyo3_path::conversion::IntoPy::into_py(self_val < other, py))
+                 },
+                #pyo3_path::pyclass::CompareOp::Le => {
+                    ::std::result::Result::Ok(#pyo3_path::conversion::IntoPy::into_py(self_val <= other, py))
+                 },
+                #pyo3_path::pyclass::CompareOp::Ge => {
+                    ::std::result::Result::Ok(#pyo3_path::conversion::IntoPy::into_py(self_val >= other, py))
+                 },
+            }
+        })
+        .unwrap_or_else(|| quote! { _ => ::std::result::Result::Ok(py.NotImplemented()) });
+
+    Ok(quote! {
+        #eq_arms
+        #ord_arms
+    })
 }
 
 fn pyclass_richcmp_simple_enum(
@@ -1723,7 +1756,7 @@ fn pyclass_richcmp_simple_enum(
         return Ok((None, None));
     }
 
-    let arms = pyclass_richcmp_arms(&options, ctx);
+    let arms = pyclass_richcmp_arms(&options, ctx)?;
 
     let eq = options.eq.map(|eq| {
         quote_spanned! { eq.span() =>
@@ -1732,7 +1765,6 @@ fn pyclass_richcmp_simple_enum(
                 let other = &*other.borrow();
                 return match op {
                     #arms
-                    _ => ::std::result::Result::Ok(py.NotImplemented())
                 }
             }
         }
@@ -1746,7 +1778,6 @@ fn pyclass_richcmp_simple_enum(
             }) {
                 return match op {
                     #arms
-                    _ => ::std::result::Result::Ok(py.NotImplemented())
                 }
             }
         }
@@ -1786,7 +1817,7 @@ fn pyclass_richcmp(
         bail_spanned!(eq_int.span() => "`eq_int` can only be used on simple enums.")
     }
 
-    let arms = pyclass_richcmp_arms(options, ctx);
+    let arms = pyclass_richcmp_arms(options, ctx)?;
     if options.eq.is_some() {
         let mut richcmp_impl = parse_quote! {
             fn __pyo3__generated____richcmp__(
@@ -1799,7 +1830,6 @@ fn pyclass_richcmp(
                 let other = &*#pyo3_path::types::PyAnyMethods::downcast::<Self>(other)?.borrow();
                 match op {
                     #arms
-                    _ => ::std::result::Result::Ok(py.NotImplemented())
                 }
             }
         };