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

This commit is contained in:
Adam Gutglick
2019-02-21 19:28:46 -05:00
parent 38b4c10833
commit f333c75fd6
3 changed files with 30 additions and 5 deletions

View File

@@ -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())))
}

View File

@@ -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())
}