diff --git a/tests/snippets/math_module.py b/tests/snippets/math_module.py index dc7fc0706..8a6aee0e9 100644 --- a/tests/snippets/math_module.py +++ b/tests/snippets/math_module.py @@ -87,3 +87,13 @@ assert math.frexp(1.5) == (0.75, 1) assert math.frexp(float('inf')) == (float('inf'), 0) assert str(math.frexp(float('nan'))) == str((float('nan'), 0)) assert_raises(TypeError, lambda: math.frexp(None)) + +assert math.gcd(0, 0) == 0 +assert math.gcd(1, 0) == 1 +assert math.gcd(0, 1) == 1 +assert math.gcd(1, 1) == 1 +assert math.gcd(-1, 1) == 1 +assert math.gcd(1, -1) == 1 +assert math.gcd(-1, -1) == 1 +assert math.gcd(125, -255) == 5 +assert_raises(TypeError, lambda: math.gcd(1.1, 2)) diff --git a/vm/src/stdlib/math.rs b/vm/src/stdlib/math.rs index 01e3fe6c1..36dad8bca 100644 --- a/vm/src/stdlib/math.rs +++ b/vm/src/stdlib/math.rs @@ -7,6 +7,7 @@ use statrs::function::erf::{erf, erfc}; use statrs::function::gamma::{gamma, ln_gamma}; use crate::function::PyFuncArgs; +use crate::obj::objint::PyIntRef; use crate::obj::{objfloat, objtype}; use crate::pyobject::{PyObjectRef, PyResult, TypeProtocol}; use crate::vm::VirtualMachine; @@ -225,6 +226,11 @@ fn math_frexp(value: PyObjectRef, vm: &VirtualMachine) -> PyResult { ) } +fn math_gcd(a: PyIntRef, b: PyIntRef, vm: &VirtualMachine) -> PyResult { + use num_integer::Integer; + Ok(vm.new_int(a.as_bigint().gcd(b.as_bigint()))) +} + pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { let ctx = &vm.ctx; @@ -279,6 +285,9 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { "ceil" => ctx.new_rustfunc(math_ceil), "floor" => ctx.new_rustfunc(math_floor), + // Gcd function + "gcd" => ctx.new_rustfunc(math_gcd), + // Constants: "pi" => ctx.new_float(std::f64::consts::PI), // 3.14159... "e" => ctx.new_float(std::f64::consts::E), // 2.71..