Merge pull request RustPython#886 from youknowone/pycomplex

 Add complex.__bool__, complex.__sub__, OverflowError/float support of complex.__add__

Author: coolreader18

tests/snippets/builtin_complex.py

@@ -1,3 +1,5 @@
+from testutils import assertRaises
+
 # __abs__
 
 assert abs(complex(3, 4)) == 5
@@ -27,6 +29,12 @@
 assert -complex(1, -1) == complex(-1, 1)
 assert -complex(0, 0) == complex(0, 0)
 
+# __bool__
+
+assert bool(complex(0, 0)) is False
+assert bool(complex(0, 1)) is True
+assert bool(complex(1, 0)) is True
+
 # real
 
 a = complex(3, 4)
@@ -44,3 +52,32 @@
 assert 1j + 1 == complex(1, 1)
 assert (1j + 1) + 3 == complex(4, 1)
 assert 3 + (1j + 1) == complex(4, 1)
+
+# float and complex addition
+assert 1.1 + 1.2j == complex(1.1, 1.2)
+assert 1.3j + 1.4 == complex(1.4, 1.3)
+assert (1.5j + 1.6) + 3 == complex(4.6, 1.5)
+assert 3.5 + (1.1j + 1.2) == complex(4.7, 1.1)
+
+# subtraction
+assert 1 - 1j == complex(1, -1)
+assert 1j - 1 == complex(-1, 1)
+assert 2j - 1j == complex(0, 1)
+
+# type error addition
+with assertRaises(TypeError):
+    assert 1j + 'str'
+with assertRaises(TypeError):
+    assert 1j - 'str'
+with assertRaises(TypeError):
+    assert 'str' + 1j
+with assertRaises(TypeError):
+    assert 'str' - 1j
+
+# overflow
+with assertRaises(OverflowError):
+    complex(10 ** 1000, 0)
+with assertRaises(OverflowError):
+    complex(0, 10 ** 1000)
+with assertRaises(OverflowError):
+    complex(0, 0) + 10 ** 1000

vm/src/obj/objcomplex.rs

@@ -1,14 +1,15 @@
 use num_complex::Complex64;
-use num_traits::ToPrimitive;
+use num_traits::{ToPrimitive, Zero};
 
 use crate::function::OptionalArg;
-use crate::pyobject::{PyContext, PyObjectRef, PyRef, PyResult, PyValue};
+use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue};
 use crate::vm::VirtualMachine;
 
 use super::objfloat::{self, PyFloat};
 use super::objint;
 use super::objtype::{self, PyClassRef};
 
+#[pyclass(name = "complex")]
 #[derive(Debug, Copy, Clone, PartialEq)]
 pub struct PyComplex {
     value: Complex64,
@@ -28,24 +29,14 @@ impl From<Complex64> for PyComplex {
 }
 
 pub fn init(context: &PyContext) {
-    let complex_type = &context.complex_type;
-
+    PyComplex::extend_class(context, &context.complex_type);
     let complex_doc =
         "Create a complex number from a real part and an optional imaginary part.\n\n\
          This is equivalent to (real + imag*1j) where imag defaults to 0.";
 
-    extend_class!(context, complex_type, {
+    extend_class!(context, &context.complex_type, {
         "__doc__" => context.new_str(complex_doc.to_string()),
-        "__abs__" => context.new_rustfunc(PyComplexRef::abs),
-        "__add__" => context.new_rustfunc(PyComplexRef::add),
-        "__eq__" => context.new_rustfunc(PyComplexRef::eq),
-        "__neg__" => context.new_rustfunc(PyComplexRef::neg),
         "__new__" => context.new_rustfunc(PyComplexRef::new),
-        "__radd__" => context.new_rustfunc(PyComplexRef::radd),
-        "__repr__" => context.new_rustfunc(PyComplexRef::repr),
-        "conjugate" => context.new_rustfunc(PyComplexRef::conjugate),
-        "imag" => context.new_property(PyComplexRef::imag),
-        "real" => context.new_property(PyComplexRef::real)
     });
 }
 
@@ -73,49 +64,87 @@ impl PyComplexRef {
         let value = Complex64::new(real, imag);
         PyComplex { value }.into_ref_with_type(vm, cls)
     }
+}
 
-    fn real(self, _vm: &VirtualMachine) -> PyFloat {
+fn to_complex(value: PyObjectRef, vm: &VirtualMachine) -> PyResult<Option<Complex64>> {
+    if objtype::isinstance(&value, &vm.ctx.int_type()) {
+        match objint::get_value(&value).to_f64() {
+            Some(v) => Ok(Some(Complex64::new(v, 0.0))),
+            None => Err(vm.new_overflow_error("int too large to convert to float".to_string())),
+        }
+    } else if objtype::isinstance(&value, &vm.ctx.float_type()) {
+        let v = objfloat::get_value(&value);
+        Ok(Some(Complex64::new(v, 0.0)))
+    } else {
+        Ok(None)
+    }
+}
+
+#[pyimpl]
+impl PyComplex {
+    #[pyproperty(name = "real")]
+    fn real(&self, _vm: &VirtualMachine) -> PyFloat {
         self.value.re.into()
     }
 
-    fn imag(self, _vm: &VirtualMachine) -> PyFloat {
+    #[pyproperty(name = "imag")]
+    fn imag(&self, _vm: &VirtualMachine) -> PyFloat {
         self.value.im.into()
     }
 
-    fn abs(self, _vm: &VirtualMachine) -> PyFloat {
+    #[pymethod(name = "__abs__")]
+    fn abs(&self, _vm: &VirtualMachine) -> PyFloat {
         let Complex64 { im, re } = self.value;
         re.hypot(im).into()
     }
 
-    fn add(self, other: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
+    #[pymethod(name = "__add__")]
+    fn add(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
         if objtype::isinstance(&other, &vm.ctx.complex_type()) {
-            vm.ctx.new_complex(self.value + get_value(&other))
-        } else if objtype::isinstance(&other, &vm.ctx.int_type()) {
-            vm.ctx.new_complex(Complex64::new(
-                self.value.re + objint::get_value(&other).to_f64().unwrap(),
-                self.value.im,
-            ))
+            Ok(vm.ctx.new_complex(self.value + get_value(&other)))
         } else {
-            vm.ctx.not_implemented()
+            self.radd(other, vm)
         }
     }
 
-    fn radd(self, other: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
-        if objtype::isinstance(&other, &vm.ctx.int_type()) {
-            vm.ctx.new_complex(Complex64::new(
-                self.value.re + objint::get_value(&other).to_f64().unwrap(),
-                self.value.im,
-            ))
+    #[pymethod(name = "__radd__")]
+    fn radd(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        match to_complex(other, vm) {
+            Ok(Some(other)) => Ok(vm.ctx.new_complex(self.value + other)),
+            Ok(None) => Ok(vm.ctx.not_implemented()),
+            Err(err) => Err(err),
+        }
+    }
+
+    #[pymethod(name = "__sub__")]
+    fn sub(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        if objtype::isinstance(&other, &vm.ctx.complex_type()) {
+            Ok(vm.ctx.new_complex(self.value - get_value(&other)))
         } else {
-            vm.ctx.not_implemented()
+            match to_complex(other, vm) {
+                Ok(Some(other)) => Ok(vm.ctx.new_complex(self.value - other)),
+                Ok(None) => Ok(vm.ctx.not_implemented()),
+                Err(err) => Err(err),
+            }
+        }
+    }
+
+    #[pymethod(name = "__rsub__")]
+    fn rsub(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        match to_complex(other, vm) {
+            Ok(Some(other)) => Ok(vm.ctx.new_complex(other - self.value)),
+            Ok(None) => Ok(vm.ctx.not_implemented()),
+            Err(err) => Err(err),
         }
     }
 
-    fn conjugate(self, _vm: &VirtualMachine) -> PyComplex {
+    #[pymethod(name = "conjugate")]
+    fn conjugate(&self, _vm: &VirtualMachine) -> PyComplex {
         self.value.conj().into()
     }
 
-    fn eq(self, other: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
+    #[pymethod(name = "__eq__")]
+    fn eq(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
         let result = if objtype::isinstance(&other, &vm.ctx.complex_type()) {
             self.value == get_value(&other)
         } else if objtype::isinstance(&other, &vm.ctx.int_type()) {
@@ -132,16 +161,23 @@ impl PyComplexRef {
         vm.ctx.new_bool(result)
     }
 
-    fn neg(self, _vm: &VirtualMachine) -> PyComplex {
+    #[pymethod(name = "__neg__")]
+    fn neg(&self, _vm: &VirtualMachine) -> PyComplex {
         PyComplex::from(-self.value)
     }
 
-    fn repr(self, _vm: &VirtualMachine) -> String {
+    #[pymethod(name = "__repr__")]
+    fn repr(&self, _vm: &VirtualMachine) -> String {
         let Complex64 { re, im } = self.value;
         if re == 0.0 {
             format!("{}j", im)
         } else {
             format!("({}+{}j)", re, im)
         }
     }
+
+    #[pymethod(name = "__bool__")]
+    fn bool(&self, _vm: &VirtualMachine) -> bool {
+        self.value != Complex64::zero()
+    }
 }