Merge pull request RustPython#891 from ntrinquier/add-methods-for-numeric-types

Add methods for the complex type

Author: windelbouwman

tests/snippets/builtin_complex.py

@@ -24,6 +24,11 @@
 assert complex(1, 2) != 'foo'
 assert complex(1, 2).__eq__('foo') == NotImplemented
 
+# __mul__
+
+assert complex(2, -3) * complex(-5, 7) == complex(11, 29)
+assert complex(2, -3) * 5 == complex(10, -15)
+
 # __neg__
 
 assert -complex(1, -1) == complex(-1, 1)

vm/src/obj/objcomplex.rs

@@ -67,7 +67,9 @@ impl PyComplexRef {
 }
 
 fn to_complex(value: PyObjectRef, vm: &VirtualMachine) -> PyResult<Option<Complex64>> {
-    if objtype::isinstance(&value, &vm.ctx.int_type()) {
+    if objtype::isinstance(&value, &vm.ctx.complex_type()) {
+        Ok(Some(get_value(&value)))
+    } else 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())),
@@ -161,6 +163,28 @@ impl PyComplex {
         vm.ctx.new_bool(result)
     }
 
+    #[pymethod(name = "__float__")]
+    fn float(&self, vm: &VirtualMachine) -> PyResult {
+        return Err(vm.new_type_error(String::from("Can't convert complex to float")));
+    }
+
+    #[pymethod(name = "__int__")]
+    fn int(&self, vm: &VirtualMachine) -> PyResult {
+        return Err(vm.new_type_error(String::from("Can't convert complex to int")));
+    }
+
+    #[pymethod(name = "__mul__")]
+    fn mul(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        match to_complex(other, vm) {
+            Ok(Some(other)) => Ok(vm.ctx.new_complex(Complex64::new(
+                self.value.re * other.re - self.value.im * other.im,
+                self.value.re * other.im + self.value.im * other.re,
+            ))),
+            Ok(None) => Ok(vm.ctx.not_implemented()),
+            Err(err) => Err(err),
+        }
+    }
+
     #[pymethod(name = "__neg__")]
     fn neg(&self, _vm: &VirtualMachine) -> PyComplex {
         PyComplex::from(-self.value)