Added the ability to do addition between complex numbers and ints.

RustPython · windelbouwman · Feb 24, 2019 · Feb 11, 2019 · Feb 11, 2019 · Feb 12, 2019
commit f333c75fd6d2ec963bc326ab89cb82ca690f00b2
diff --git a/tests/snippets/builtin_complex.py b/tests/snippets/builtin_complex.py
@@ -42,3 +42,5 @@
 # int and complex addition
 assert 1 + 1j == complex(1, 1)
 assert 1j + 1 == complex(1, 1)
+assert (1j + 1) + 3 == complex(4, 1)
+assert 3 + (1j + 1) == complex(4, 1)
diff --git a/vm/src/obj/objcomplex.rs b/vm/src/obj/objcomplex.rs
@@ -17,6 +17,11 @@ pub fn init(context: &PyContext) {
 
     context.set_attr(&complex_type, "__abs__", context.new_rustfunc(complex_abs));
     context.set_attr(&complex_type, "__add__", context.new_rustfunc(complex_add));
+    context.set_attr(
+        &complex_type,
+        "__radd__",
+        context.new_rustfunc(complex_radd),
+    );
     context.set_attr(&complex_type, "__eq__", context.new_rustfunc(complex_eq));
     context.set_attr(&complex_type, "__neg__", context.new_rustfunc(complex_neg));
     context.set_attr(&complex_type, "__new__", context.new_rustfunc(complex_new));
@@ -107,7 +112,29 @@ fn complex_add(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     if objtype::isinstance(i2, &vm.ctx.complex_type()) {
         Ok(vm.ctx.new_complex(v1 + get_value(i2)))
     } else if objtype::isinstance(i2, &vm.ctx.int_type()) {
-        Ok(vm.ctx.new_complex(Complex64::new(v1.re + objint::get_value(i2).to_f64().unwrap(), v1.im)))
+        Ok(vm.ctx.new_complex(Complex64::new(
+            v1.re + objint::get_value(i2).to_f64().unwrap(),
+            v1.im,
+        )))
+    } else {
+        Err(vm.new_type_error(format!("Cannot add {} and {}", i.borrow(), i2.borrow())))
+    }
+}
+
+fn complex_radd(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(i, Some(vm.ctx.complex_type())), (i2, None)]
+    );
+
+    let v1 = get_value(i);
+
+    if objtype::isinstance(i2, &vm.ctx.int_type()) {
+        Ok(vm.ctx.new_complex(Complex64::new(
+            v1.re + objint::get_value(i2).to_f64().unwrap(),
+            v1.im,
+        )))
     } else {
         Err(vm.new_type_error(format!("Cannot add {} and {}", i.borrow(), i2.borrow())))
     }

diff --git a/vm/src/obj/objint.rs b/vm/src/obj/objint.rs
@@ -1,7 +1,6 @@
 use super::objfloat;
 use super::objstr;
 use super::objtype;
-use super::objcomplex;
 use crate::format::FormatSpec;
 use crate::pyobject::{
     FromPyObjectRef, PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectRef, PyResult,
@@ -11,7 +10,6 @@ use crate::vm::VirtualMachine;
 use num_bigint::{BigInt, ToBigInt};
 use num_integer::Integer;
 use num_traits::{Pow, Signed, ToPrimitive, Zero};
-use num_complex::Complex64;
 use std::hash::{Hash, Hasher};
 
 // This proxy allows for easy switching between types.
@@ -291,8 +289,6 @@ fn int_add(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     );
     if objtype::isinstance(other, &vm.ctx.int_type()) {
         Ok(vm.ctx.new_int(get_value(zelf) + get_value(other)))
-    } else if objtype::isinstance(other, &vm.ctx.complex_type()) {
-        Ok(vm.ctx.new_complex(Complex64::new(get_value(zelf).to_f64().unwrap() + objcomplex::get_value(other).re, objcomplex::get_value(other).im)))
     } else {
         Ok(vm.ctx.not_implemented())
     }