Add math.ceil / math.floor

Author: youknowone

tests/snippets/math_basics.py

@@ -1,4 +1,3 @@
-import math
 from testutils import assertRaises
 
 a = 4
@@ -19,34 +18,3 @@
 assert a - 3 == 1
 assert -a == -4
 assert +a == 4
-
-# 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 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'

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,7 +7,7 @@ use statrs::function::erf::{erf, erfc};
 use statrs::function::gamma::{gamma, ln_gamma};
 
 use crate::function::PyFuncArgs;
-use crate::obj::objfloat;
+use crate::obj::{objfloat, objtype};
 use crate::pyobject::{PyObjectRef, PyResult, TypeProtocol};
 use crate::vm::VirtualMachine;
 
@@ -172,20 +172,43 @@ 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) {
+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,
-            MAGIC_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;
 
@@ -235,6 +258,8 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
 
         // 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...