Implement Add/Sub/Mul/Div for &PyComplex

Cargo.toml

@@ -27,6 +27,7 @@ pyo3cls = { path = "pyo3cls", version = "=0.5.0-alpha.1" }
 mashup = "0.1.7"
 
 [dev-dependencies]
+assert_approx_eq = "1.0"
 docmatic = "0.1.2"
 
 [build-dependencies]

src/ffi3/complexobject.rs

@@ -22,3 +22,33 @@ extern "C" {
     pub fn PyComplex_RealAsDouble(op: *mut PyObject) -> c_double;
     pub fn PyComplex_ImagAsDouble(op: *mut PyObject) -> c_double;
 }
+
+#[cfg(not(Py_LIMITED_API))]
+#[repr(C)]
+#[derive(Copy, Clone)]
+pub struct Py_complex {
+    pub real: c_double,
+    pub imag: c_double,
+}
+
+#[cfg(not(Py_LIMITED_API))]
+#[repr(C)]
+#[derive(Copy, Clone)]
+pub struct PyComplexObject {
+    _ob_base: PyObject,
+    pub cval: Py_complex,
+}
+
+#[cfg(not(Py_LIMITED_API))]
+#[cfg_attr(windows, link(name = "pythonXY"))]
+extern "C" {
+    pub fn _Py_c_sum(left: Py_complex, right: Py_complex) -> Py_complex;
+    pub fn _Py_c_diff(left: Py_complex, right: Py_complex) -> Py_complex;
+    pub fn _Py_c_neg(complex: Py_complex) -> Py_complex;
+    pub fn _Py_c_prod(left: Py_complex, right: Py_complex) -> Py_complex;
+    pub fn _Py_c_quot(dividend: Py_complex, divisor: Py_complex) -> Py_complex;
+    pub fn _Py_c_pow(num: Py_complex, exp: Py_complex) -> Py_complex;
+    pub fn _Py_c_abs(arg: Py_complex) -> c_double;
+    pub fn PyComplex_FromCComplex(v: Py_complex) -> *mut PyObject;
+    pub fn PyComplex_AsCComplex(op: *mut PyObject) -> Py_complex;
+}

src/lib.rs

@@ -128,6 +128,9 @@ extern crate spin;
 // We need that reexport for wrap_function
 #[doc(hidden)]
 pub extern crate mashup;
+#[cfg(test)]
+#[macro_use]
+extern crate assert_approx_eq;
 
 /// Rust FFI declarations for Python
 pub mod ffi;

src/objects/complex.rs

@@ -1,6 +1,9 @@
 use ffi;
+use instance::PyObjectWithToken;
 use object::PyObject;
 use python::{Python, ToPyPointer};
+#[cfg(any(not(Py_LIMITED_API), not(Py_3)))]
+use std::ops::*;
 use std::os::raw::c_double;
 
 /// Represents a Python `complex`.
@@ -24,6 +27,62 @@ impl PyComplex {
     }
 }
 
+#[cfg(any(not(Py_LIMITED_API), not(Py_3)))]
+#[inline(always)]
+unsafe fn complex_operation(
+    l: &PyComplex,
+    r: &PyComplex,
+    operation: unsafe extern "C" fn(ffi::Py_complex, ffi::Py_complex) -> ffi::Py_complex,
+) -> *mut ffi::PyObject {
+    let l_val = (*(l.as_ptr() as *mut ffi::PyComplexObject)).cval;
+    let r_val = (*(r.as_ptr() as *mut ffi::PyComplexObject)).cval;
+    ffi::PyComplex_FromCComplex(operation(l_val, r_val))
+}
+
+#[cfg(any(not(Py_LIMITED_API), not(Py_3)))]
+impl<'py> Add for &'py PyComplex {
+    type Output = &'py PyComplex;
+    fn add(self, other: &'py PyComplex) -> &'py PyComplex {
+        unsafe {
+            self.py()
+                .from_owned_ptr(complex_operation(self, other, ffi::_Py_c_sum))
+        }
+    }
+}
+
+#[cfg(any(not(Py_LIMITED_API), not(Py_3)))]
+impl<'py> Sub for &'py PyComplex {
+    type Output = &'py PyComplex;
+    fn sub(self, other: &'py PyComplex) -> &'py PyComplex {
+        unsafe {
+            self.py()
+                .from_owned_ptr(complex_operation(self, other, ffi::_Py_c_diff))
+        }
+    }
+}
+
+#[cfg(any(not(Py_LIMITED_API), not(Py_3)))]
+impl<'py> Mul for &'py PyComplex {
+    type Output = &'py PyComplex;
+    fn mul(self, other: &'py PyComplex) -> &'py PyComplex {
+        unsafe {
+            self.py()
+                .from_owned_ptr(complex_operation(self, other, ffi::_Py_c_prod))
+        }
+    }
+}
+
+#[cfg(any(not(Py_LIMITED_API), not(Py_3)))]
+impl<'py> Div for &'py PyComplex {
+    type Output = &'py PyComplex;
+    fn div(self, other: &'py PyComplex) -> &'py PyComplex {
+        unsafe {
+            self.py()
+                .from_owned_ptr(complex_operation(self, other, ffi::_Py_c_quot))
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::PyComplex;
@@ -36,4 +95,52 @@ mod test {
         assert_eq!(complex.real(), 3.0);
         assert_eq!(complex.imag(), 1.2);
     }
+
+    #[cfg(any(not(Py_LIMITED_API), not(Py_3)))]
+    #[test]
+    fn test_add() {
+        let gil = Python::acquire_gil();
+        let py = gil.python();
+        let l = PyComplex::from_doubles(py, 3.0, 1.2);
+        let r = PyComplex::from_doubles(py, 1.0, 2.6);
+        let res = l + r;
+        assert_approx_eq!(res.real(), 4.0);
+        assert_approx_eq!(res.imag(), 3.8);
+    }
+
+    #[cfg(any(not(Py_LIMITED_API), not(Py_3)))]
+    #[test]
+    fn test_sub() {
+        let gil = Python::acquire_gil();
+        let py = gil.python();
+        let l = PyComplex::from_doubles(py, 3.0, 1.2);
+        let r = PyComplex::from_doubles(py, 1.0, 2.6);
+        let res = l - r;
+        assert_approx_eq!(res.real(), 2.0);
+        assert_approx_eq!(res.imag(), -1.4);
+    }
+
+    #[cfg(any(not(Py_LIMITED_API), not(Py_3)))]
+    #[test]
+    fn test_mul() {
+        let gil = Python::acquire_gil();
+        let py = gil.python();
+        let l = PyComplex::from_doubles(py, 3.0, 1.2);
+        let r = PyComplex::from_doubles(py, 1.0, 2.6);
+        let res = l * r;
+        assert_approx_eq!(res.real(), -0.12);
+        assert_approx_eq!(res.imag(), 9.0);
+    }
+
+    #[cfg(any(not(Py_LIMITED_API), not(Py_3)))]
+    #[test]
+    fn test_div() {
+        let gil = Python::acquire_gil();
+        let py = gil.python();
+        let l = PyComplex::from_doubles(py, 3.0, 1.2);
+        let r = PyComplex::from_doubles(py, 1.0, 2.6);
+        let res = l / r;
+        assert_approx_eq!(res.real(), 0.7886597938144329);
+        assert_approx_eq!(res.imag(), -0.8505154639175257);
+    }
 }