Add math.trunc

Author: youknowone

tests/snippets/math_basics.py

@@ -1,3 +1,6 @@
+import math
+from testutils import assertRaises
+
 a = 4
 
 #print(a ** 3)
@@ -17,7 +20,6 @@
 assert -a == -4
 assert +a == 4
 
-# import math
 # assert(math.exp(2) == math.exp(2.0))
 # assert(math.exp(True) == math.exp(1.0))
 #
@@ -27,3 +29,24 @@
 #         return 1.1111
 #
 # assert math.log(1.1111) == math.log(Conversible())
+
+# roundings
+assert math.trunc(1) == 1
+
+class A(object):
+    def __trunc__(self):
+        return 2
+
+assert math.trunc(A()) == 2
+
+class A(object):
+    def __trunc__(self):
+        return 2.0
+
+assert math.trunc(A()) == 2.0
+
+class A(object):
+    def __trunc__(self):
+        return 'str'
+
+assert math.trunc(A()) == 'str'

vm/src/stdlib/math.rs

@@ -8,7 +8,7 @@ use statrs::function::gamma::{gamma, ln_gamma};
 
 use crate::function::PyFuncArgs;
 use crate::obj::objfloat;
-use crate::pyobject::{PyObjectRef, PyResult};
+use crate::pyobject::{PyObjectRef, PyResult, TypeProtocol};
 use crate::vm::VirtualMachine;
 
 // Helper macro:
@@ -172,6 +172,20 @@ fn math_lgamma(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
     }
 }
 
+fn math_trunc(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(value, None)]);
+    const MAGIC_NAME: &str = "__trunc__";
+    if let Ok(method) = vm.get_method(value.clone(), MAGIC_NAME) {
+        vm.invoke(method, vec![])
+    } else {
+        Err(vm.new_type_error(format!(
+            "TypeError: type {} doesn't define {} method",
+            value.class().name,
+            MAGIC_NAME,
+        )))
+    }
+}
+
 pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     let ctx = &vm.ctx;
 
@@ -219,6 +233,9 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
         "gamma" => ctx.new_rustfunc(math_gamma),
         "lgamma" => ctx.new_rustfunc(math_lgamma),
 
+        // Rounding functions:
+        "trunc" => ctx.new_rustfunc(math_trunc),
+
         // Constants:
         "pi" => ctx.new_float(std::f64::consts::PI), // 3.14159...
         "e" => ctx.new_float(std::f64::consts::E), // 2.71..