Merge pull request #892 from youknowone/float-round

Add basic float.__round__ without ndigits support

tests/snippets/floats.py

@@ -133,7 +133,17 @@ assert 1.2.__float__() == 1.2
 assert 1.2.__trunc__() == 1
 assert int(1.2) == 1
 assert float(1.2) == 1.2
-# assert math.trunc(1.2) == 1
+assert math.trunc(1.2) == 1
+assert_raises(OverflowError, float('inf').__trunc__)
+assert_raises(ValueError, float('nan').__trunc__)
+assert 0.5.__round__() == 0.0
+assert 1.5.__round__() == 2.0
+assert 0.5.__round__(0) == 0.0
+assert 1.5.__round__(0) == 2.0
+assert 0.5.__round__(None) == 0.0
+assert 1.5.__round__(None) == 2.0
+assert_raises(OverflowError, float('inf').__round__)
+assert_raises(ValueError, float('nan').__round__)
 
 assert 1.2 ** 2 == 1.44
 assert_raises(OverflowError, lambda: 1.2 ** (10 ** 1000))

tests/snippets/math_basics.py

@@ -1,4 +1,4 @@
-from testutils import assertRaises
+from testutils import assert_raises
 
 a = 4
 
@@ -18,3 +18,23 @@ assert 2 == a / 2
 assert a - 3 == 1
 assert -a == -4
 assert +a == 4
+
+assert round(1.2) == 1
+assert round(1.8) == 2
+assert round(0.5) == 0
+assert round(1.5) == 2
+assert round(-0.5) == 0
+assert round(-1.5) == -2
+
+assert_raises(
+    ValueError,
+    lambda: round(float('nan')),
+    'ValueError: cannot convert float NaN to integer')
+assert_raises(
+    OverflowError,
+    lambda: round(float('inf')),
+    'OverflowError: cannot convert float NaN to integer')
+assert_raises(
+    OverflowError,
+    lambda: round(-float('inf')),
+    'OverflowError: cannot convert float NaN to integer')

vm/src/obj/objfloat.rs

@@ -2,14 +2,16 @@ use super::objbytes;
 use super::objint;
 use super::objstr;
 use super::objtype;
+use crate::function::OptionalArg;
 use crate::obj::objtype::PyClassRef;
 use crate::pyobject::{
-    IntoPyObject, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol,
+    IdProtocol, IntoPyObject, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue,
+    TypeProtocol,
 };
 use crate::vm::VirtualMachine;
 use num_bigint::{BigInt, ToBigInt};
 use num_rational::Ratio;
-use num_traits::ToPrimitive;
+use num_traits::{ToPrimitive, Zero};
 
 #[pyclass(name = "float")]
 #[derive(Debug, Copy, Clone, PartialEq)]
@@ -67,6 +69,28 @@ fn inner_mod(v1: f64, v2: f64, vm: &VirtualMachine) -> PyResult<f64> {
     }
 }
 
+fn try_to_bigint(value: f64, vm: &VirtualMachine) -> PyResult<BigInt> {
+    match value.to_bigint() {
+        Some(int) => Ok(int),
+        None => {
+            if value.is_infinite() {
+                Err(vm.new_overflow_error(
+                    "OverflowError: cannot convert float NaN to integer".to_string(),
+                ))
+            } else if value.is_nan() {
+                Err(vm
+                    .new_value_error("ValueError: cannot convert float NaN to integer".to_string()))
+            } else {
+                // unreachable unless BigInt has a bug
+                unreachable!(
+                    "A finite float value failed to be converted to bigint: {}",
+                    value
+                )
+            }
+        }
+    }
+}
+
 fn inner_floordiv(v1: f64, v2: f64, vm: &VirtualMachine) -> PyResult<f64> {
     if v2 != 0.0 {
         Ok((v1 / v2).floor())
@@ -343,12 +367,54 @@ impl PyFloat {
     }
 
     #[pymethod(name = "__trunc__")]
-    fn trunc(&self, _vm: &VirtualMachine) -> BigInt {
-        self.value.to_bigint().unwrap()
+    fn trunc(&self, vm: &VirtualMachine) -> PyResult<BigInt> {
+        try_to_bigint(self.value, vm)
+    }
+
+    #[pymethod(name = "__round__")]
+    fn round(&self, ndigits: OptionalArg<PyObjectRef>, vm: &VirtualMachine) -> PyResult {
+        let ndigits = match ndigits {
+            OptionalArg::Missing => None,
+            OptionalArg::Present(ref value) => {
+                if !vm.get_none().is(value) {
+                    let ndigits = if objtype::isinstance(value, &vm.ctx.int_type()) {
+                        objint::get_value(value)
+                    } else {
+                        return Err(vm.new_type_error(format!(
+                            "TypeError: '{}' object cannot be interpreted as an integer",
+                            value.class().name
+                        )));
+                    };
+                    if ndigits.is_zero() {
+                        None
+                    } else {
+                        Some(ndigits)
+                    }
+                } else {
+                    None
+                }
+            }
+        };
+        if ndigits.is_none() {
+            let fract = self.value.fract();
+            let value = if fract.abs() == 0.5 {
+                if self.value.trunc() % 2.0 == 0.0 {
+                    self.value - fract
+                } else {
+                    self.value + fract
+                }
+            } else {
+                self.value.round()
+            };
+            let int = try_to_bigint(value, vm)?;
+            Ok(vm.ctx.new_int(int))
+        } else {
+            Ok(vm.ctx.not_implemented())
+        }
     }
 
     #[pymethod(name = "__int__")]
-    fn int(&self, vm: &VirtualMachine) -> BigInt {
+    fn int(&self, vm: &VirtualMachine) -> PyResult<BigInt> {
         self.trunc(vm)
     }