From fc863aaba5466e8c109fa3f8e9b2209aaeaa3e11 Mon Sep 17 00:00:00 2001 From: Joey Hain Date: Fri, 8 Feb 2019 18:20:55 -0800 Subject: [PATCH] Add complex.__abs__ --- tests/snippets/builtin_complex.py | 3 +++ vm/src/obj/objcomplex.rs | 8 ++++++++ 2 files changed, 11 insertions(+) create mode 100644 tests/snippets/builtin_complex.py diff --git a/tests/snippets/builtin_complex.py b/tests/snippets/builtin_complex.py new file mode 100644 index 000000000..5ac70ff59 --- /dev/null +++ b/tests/snippets/builtin_complex.py @@ -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 diff --git a/vm/src/obj/objcomplex.rs b/vm/src/obj/objcomplex.rs index 353b89b7f..00556096b 100644 --- a/vm/src/obj/objcomplex.rs +++ b/vm/src/obj/objcomplex.rs @@ -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,