From 38b4c10833810ffc47fa78af95edfefe8dfe99eb Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Thu, 21 Feb 2019 19:28:46 -0500 Subject: [PATCH 1/2] Added the ability to do addition between complex numbers and ints. --- tests/snippets/builtin_complex.py | 4 ++++ vm/src/obj/objcomplex.rs | 2 ++ vm/src/obj/objint.rs | 4 ++++ 3 files changed, 10 insertions(+) diff --git a/tests/snippets/builtin_complex.py b/tests/snippets/builtin_complex.py index 6447c1f5b8..b979637d76 100644 --- 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 index a7b1aaac96..ecdbbdb082 100644 --- 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 index b5e38d8303..935980de1a 100644 --- 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()) } From f333c75fd6d2ec963bc326ab89cb82ca690f00b2 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Thu, 21 Feb 2019 19:28:46 -0500 Subject: [PATCH 2/2] Added the ability to do addition between complex numbers and ints. --- tests/snippets/builtin_complex.py | 2 ++ vm/src/obj/objcomplex.rs | 29 ++++++++++++++++++++++++++++- vm/src/obj/objint.rs | 4 ---- 3 files changed, 30 insertions(+), 5 deletions(-) diff --git a/tests/snippets/builtin_complex.py b/tests/snippets/builtin_complex.py index b979637d76..c073f74337 100644 --- 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 index ecdbbdb082..9182f7bc2e 100644 --- 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 index 935980de1a..b5e38d8303 100644 --- 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()) }