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

AdamGS · AdamGS · commit 38b4c1083381 · 2019-02-21T19:28:46.000-05:00
diff --git a/tests/snippets/builtin_complex.py b/tests/snippets/builtin_complex.py
@@ -38,3 +38,7 @@
 
 assert a.imag == 4
 assert b.imag == 4
+
+# int and complex addition
+assert 1 + 1j == complex(1, 1)
+assert 1j + 1 == complex(1, 1)
diff --git a/vm/src/obj/objcomplex.rs b/vm/src/obj/objcomplex.rs
@@ -106,6 +106,8 @@ fn complex_add(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     let v1 = get_value(i);
     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)))
     } 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,6 +1,7 @@
 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,
@@ -10,6 +11,7 @@ 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.
@@ -289,6 +291,8 @@ 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())
     }

Original file line number	Diff line number	Diff line change
`@@ -106,6 +106,8 @@ fn complex_add(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {`
`106`	`106`	`let v1 = get_value(i);`
`107`	`107`	`if objtype::isinstance(i2, &vm.ctx.complex_type()) {`
`108`	`108`	`Ok(vm.ctx.new_complex(v1 + get_value(i2)))`
	`109`	`+ } else if objtype::isinstance(i2, &vm.ctx.int_type()) {`
	`110`	`+ Ok(vm.ctx.new_complex(Complex64::new(v1.re + objint::get_value(i2).to_f64().unwrap(), v1.im)))`
`109`	`111`	`} else {`
`110`	`112`	`Err(vm.new_type_error(format!("Cannot add {} and {}", i.borrow(), i2.borrow())))`
`111`	`113`	`}`