Add complex.__abs__

This commit is contained in:
Joey Hain
2019-02-08 18:20:55 -08:00
parent f454bf36d1
commit fc863aaba5
2 changed files with 11 additions and 0 deletions

View File

@@ -0,0 +1,3 @@
assert complex(3, 4).__abs__() == 5
assert complex(3, -4).__abs__() == 5
assert complex(1.5, 2.5).__abs__() == 2.9154759474226504

View File

@@ -13,6 +13,7 @@ pub fn init(context: &PyContext) {
"Create a complex number from a real part and an optional imaginary part.\n\n\
This is equivalent to (real + imag*1j) where imag defaults to 0.";
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, "__new__", context.new_rustfunc(complex_new));
context.set_attr(
@@ -70,6 +71,13 @@ fn complex_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
))
}
fn complex_abs(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
arg_check!(vm, args, required = [(zelf, Some(vm.ctx.complex_type()))]);
let Complex64 { re, im } = get_value(zelf);
Ok(vm.ctx.new_float(re.hypot(im)))
}
fn complex_add(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
arg_check!(
vm,