Skip to content

Commit 38b4c10

Browse files
committed
Added the ability to do addition between complex numbers and ints.
1 parent 5bd2db8 commit 38b4c10

File tree

3 files changed

+10
-0
lines changed

3 files changed

+10
-0
lines changed

tests/snippets/builtin_complex.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,3 +38,7 @@
3838

3939
assert a.imag == 4
4040
assert b.imag == 4
41+
42+
# int and complex addition
43+
assert 1 + 1j == complex(1, 1)
44+
assert 1j + 1 == complex(1, 1)

vm/src/obj/objcomplex.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,8 @@ fn complex_add(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
106106
let v1 = get_value(i);
107107
if objtype::isinstance(i2, &vm.ctx.complex_type()) {
108108
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)))
109111
} else {
110112
Err(vm.new_type_error(format!("Cannot add {} and {}", i.borrow(), i2.borrow())))
111113
}

vm/src/obj/objint.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use super::objfloat;
22
use super::objstr;
33
use super::objtype;
4+
use super::objcomplex;
45
use crate::format::FormatSpec;
56
use crate::pyobject::{
67
FromPyObjectRef, PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectRef, PyResult,
@@ -10,6 +11,7 @@ use crate::vm::VirtualMachine;
1011
use num_bigint::{BigInt, ToBigInt};
1112
use num_integer::Integer;
1213
use num_traits::{Pow, Signed, ToPrimitive, Zero};
14+
use num_complex::Complex64;
1315
use std::hash::{Hash, Hasher};
1416

1517
// This proxy allows for easy switching between types.
@@ -289,6 +291,8 @@ fn int_add(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
289291
);
290292
if objtype::isinstance(other, &vm.ctx.int_type()) {
291293
Ok(vm.ctx.new_int(get_value(zelf) + get_value(other)))
294+
} else if objtype::isinstance(other, &vm.ctx.complex_type()) {
295+
Ok(vm.ctx.new_complex(Complex64::new(get_value(zelf).to_f64().unwrap() + objcomplex::get_value(other).re, objcomplex::get_value(other).im)))
292296
} else {
293297
Ok(vm.ctx.not_implemented())
294298
}

0 commit comments

Comments
 (0)