Support compare() in Python 3; and add rich_compare().

src/objectprotocol.rs

@@ -70,13 +70,26 @@ pub trait ObjectProtocol : PythonObject {
     }
 
     /// Compares two Python objects.
-    /// This is equivalent to the Python expression 'cmp(self, other)'.
-    #[cfg(feature="python27-sys")]
+    ///
+    /// On Python 2, this is equivalent to the Python expression 'cmp(self, other)'.
+    ///
+    /// On Python 3, this is equivalent to:
+    /// ```
+    /// if self == other:
+    ///     return Equal
+    /// elif a < b:
+    ///     return Less
+    /// elif a > b:
+    ///     return Greater
+    /// else:
+    ///     raise TypeError("ObjectProtocol::compare(): All comparisons returned false")
+    /// ```
     fn compare<O>(&self, py: Python, other: O) -> PyResult<Ordering> where O: ToPyObject {
-        other.with_borrowed_ptr(py, |other| unsafe {
+        #[cfg(feature="python27-sys")]
+        unsafe fn do_compare(py: Python, a: *mut ffi::PyObject, b: *mut ffi::PyObject) -> PyResult<Ordering> {
             let mut result = -1;
             try!(err::error_on_minusone(py,
-                ffi::PyObject_Cmp(self.as_ptr(), other, &mut result)));
+                ffi::PyObject_Cmp(a, b, &mut result)));
             Ok(if result < 0 {
                 Ordering::Less
             } else if result > 0 {
@@ -84,6 +97,48 @@ pub trait ObjectProtocol : PythonObject {
             } else {
                 Ordering::Equal
             })
+        }
+
+        #[cfg(feature="python3-sys")]
+        unsafe fn do_compare(py: Python, a: *mut ffi::PyObject, b: *mut ffi::PyObject) -> PyResult<Ordering> {
+            let result = ffi::PyObject_RichCompareBool(a, b, ffi::Py_EQ);
+            if result == 1 {
+                return Ok(Ordering::Equal);
+            } else if result < 0 {
+                return Err(PyErr::fetch(py));
+            }
+            let result = ffi::PyObject_RichCompareBool(a, b, ffi::Py_LT);
+            if result == 1 {
+                return Ok(Ordering::Less);
+            } else if result < 0 {
+                return Err(PyErr::fetch(py));
+            }
+            let result = ffi::PyObject_RichCompareBool(a, b, ffi::Py_GT);
+            if result == 1 {
+                return Ok(Ordering::Greater);
+            } else if result < 0 {
+                return Err(PyErr::fetch(py));
+            }
+            return Err(PyErr::new::<::exc::TypeError, _>(py, "ObjectProtocol::compare(): All comparisons returned false"));
+        }
+
+        other.with_borrowed_ptr(py, |other| unsafe {
+            do_compare(py, self.as_ptr(), other)
+        })
+    }
+
+    /// Compares two Python objects.
+    ///
+    /// Depending on the value of `compare_op`, equivalent to one of the following Python expressions:
+    ///   * CompareOp::Eq: `self == other`
+    ///   * CompareOp::Ne: `self != other`
+    ///   * CompareOp::Lt: `self < other`
+    ///   * CompareOp::Le: `self <= other`
+    ///   * CompareOp::Gt: `self > other`
+    ///   * CompareOp::Ge: `self >= other`
+    fn rich_compare<O>(&self, py: Python, other: O, compare_op: ::CompareOp) -> PyResult<PyObject> where O: ToPyObject {
+        other.with_borrowed_ptr(py, |other| unsafe {
+            err::result_cast_from_owned_ptr(py, ffi::PyObject_RichCompare(self.as_ptr(), other, compare_op as libc::c_int))
         })
     }
 
@@ -249,6 +304,7 @@ mod test {
     use python::{Python, PythonObject};
     use conversion::ToPyObject;
     use objects::{PyList, PyTuple};
+    use super::ObjectProtocol;
 
     #[test]
     fn test_debug_string() {
@@ -265,5 +321,16 @@ mod test {
         let v = "Hello\n".to_py_object(py).into_object();
         assert_eq!(format!("{}", v), "Hello\n");
     }
+
+    #[test]
+    fn test_compare() {
+        use std::cmp::Ordering;
+        let gil = Python::acquire_gil();
+        let py = gil.python();
+        let one = 1i32.to_py_object(py).into_object();
+        assert_eq!(one.compare(py, 1).unwrap(), Ordering::Equal);
+        assert_eq!(one.compare(py, 2).unwrap(), Ordering::Less);
+        assert_eq!(one.compare(py, 0).unwrap(), Ordering::Greater);
+    }
 }
 

src/py_class/mod.rs

@@ -32,6 +32,7 @@ use objects::{PyObject, PyType, PyModule};
 use err::{self, PyResult};
 use ffi;
 
+// TODO: consider moving CompareOp to a different module, so that it isn't exported via two paths
 #[derive(Debug)]
 pub enum CompareOp {
     Lt = ffi::Py_LT as isize,

src/py_class/py_class.rs

@@ -280,12 +280,15 @@ py_class!(class MyIterator |py| {
 
 ## Comparison operators
 
-  * `def __richcmp__(&self, other: impl ToPyObject, op: CompareOp) -> PyResult<impl ToPyObject>`
+  * `def __richcmp__(&self, other: impl FromPyObject, op: CompareOp) -> PyResult<impl ToPyObject>`
 
-    Overloads Python comparison operations (`==`, `!=`, `<`, `<=`, `>`, and `>=`). The `op`
-    argument indicates the comparison operation being performed.
+    Overloads Python comparison operations (`==`, `!=`, `<`, `<=`, `>`, and `>=`).
+    The `op` argument indicates the comparison operation being performed.
     The return type will normally be `PyResult<bool>`, but any Python object can be returned.
 
+    If `other` is not of the type specified in the signature, the generated code will
+    automatically `return NotImplemented`.
+
   * `def __hash__(&self) -> PyResult<impl PrimInt>`
 
     Objects that compare equal must have the same hash value.