Merge pull request RustPython#888 from youknowone/math-rounding

Add trunc, ceil, floor to math module

Author: windelbouwman

tests/snippets/math_basics.py

@@ -1,3 +1,5 @@
+from testutils import assertRaises
+
 a = 4
 
 #print(a ** 3)
@@ -16,14 +18,3 @@
 assert a - 3 == 1
 assert -a == -4
 assert +a == 4
-
-# import math
-# assert(math.exp(2) == math.exp(2.0))
-# assert(math.exp(True) == math.exp(1.0))
-#
-# class Conversible():
-#     def __float__(self):
-#         print("Converting to float now!")
-#         return 1.1111
-#
-# assert math.log(1.1111) == math.log(Conversible())

tests/snippets/math_module.py

@@ -0,0 +1,80 @@
+import math
+from testutils import assertRaises
+
+# assert(math.exp(2) == math.exp(2.0))
+# assert(math.exp(True) == math.exp(1.0))
+#
+# class Conversible():
+#     def __float__(self):
+#         print("Converting to float now!")
+#         return 1.1111
+#
+# assert math.log(1.1111) == math.log(Conversible())
+
+# roundings
+assert int.__trunc__
+assert int.__floor__
+assert int.__ceil__
+
+# assert float.__trunc__
+with assertRaises(AttributeError):
+    assert float.__floor__
+with assertRaises(AttributeError):
+    assert float.__ceil__
+
+assert math.trunc(2) == 2
+assert math.ceil(3) == 3
+assert math.floor(4) == 4
+
+assert math.trunc(2.2) == 2
+assert math.ceil(3.3) == 4
+assert math.floor(4.4) == 4
+
+class A(object):
+    def __trunc__(self):
+        return 2
+
+    def __ceil__(self):
+        return 3
+
+    def __floor__(self):
+        return 4
+
+assert math.trunc(A()) == 2
+assert math.ceil(A()) == 3
+assert math.floor(A()) == 4
+
+class A(object):
+    def __trunc__(self):
+        return 2.2
+
+    def __ceil__(self):
+        return 3.3
+
+    def __floor__(self):
+        return 4.4
+
+assert math.trunc(A()) == 2.2
+assert math.ceil(A()) == 3.3
+assert math.floor(A()) == 4.4
+
+class A(object):
+    def __trunc__(self):
+        return 'trunc'
+
+    def __ceil__(self):
+        return 'ceil'
+
+    def __floor__(self):
+        return 'floor'
+
+assert math.trunc(A()) == 'trunc'
+assert math.ceil(A()) == 'ceil'
+assert math.floor(A()) == 'floor'
+
+with assertRaises(TypeError):
+    math.trunc(object())
+with assertRaises(TypeError):
+    math.ceil(object())
+with assertRaises(TypeError):
+    math.floor(object())

vm/src/stdlib/math.rs

@@ -7,8 +7,8 @@ use statrs::function::erf::{erf, erfc};
 use statrs::function::gamma::{gamma, ln_gamma};
 
 use crate::function::PyFuncArgs;
-use crate::obj::objfloat;
-use crate::pyobject::{PyObjectRef, PyResult};
+use crate::obj::{objfloat, objtype};
+use crate::pyobject::{PyObjectRef, PyResult, TypeProtocol};
 use crate::vm::VirtualMachine;
 
 // Helper macro:
@@ -172,6 +172,43 @@ fn math_lgamma(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
     }
 }
 
+fn try_magic_method(func_name: &str, vm: &VirtualMachine, value: &PyObjectRef) -> PyResult {
+    if let Ok(method) = vm.get_method(value.clone(), func_name) {
+        vm.invoke(method, vec![])
+    } else {
+        Err(vm.new_type_error(format!(
+            "TypeError: type {} doesn't define {} method",
+            value.class().name,
+            func_name,
+        )))
+    }
+}
+
+fn math_trunc(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(value, None)]);
+    try_magic_method("__trunc__", vm, value)
+}
+
+fn math_ceil(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(value, None)]);
+    if objtype::isinstance(value, &vm.ctx.float_type) {
+        let v = objfloat::get_value(value);
+        Ok(vm.ctx.new_float(v.ceil()))
+    } else {
+        try_magic_method("__ceil__", vm, value)
+    }
+}
+
+fn math_floor(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(value, None)]);
+    if objtype::isinstance(value, &vm.ctx.float_type) {
+        let v = objfloat::get_value(value);
+        Ok(vm.ctx.new_float(v.floor()))
+    } else {
+        try_magic_method("__floor__", vm, value)
+    }
+}
+
 pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     let ctx = &vm.ctx;
 
@@ -219,6 +256,11 @@ 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),
+        "ceil" => ctx.new_rustfunc(math_ceil),
+        "floor" => ctx.new_rustfunc(math_floor),
+
         // Constants:
         "pi" => ctx.new_float(std::f64::consts::PI), // 3.14159...
         "e" => ctx.new_float(std::f64::consts::E), // 2.71..