py/modbuiltins: Add support for rounding integers.

honza-klu · dpgeorge · commit b318ebf1015c · 2018-05-22T14:18:16.000+10:00
As per CPython semantics.  This feature is controlled by
MICROPY_PY_BUILTINS_ROUND_INT which is disabled by default.
diff --git a/py/modbuiltins.c b/py/modbuiltins.c
@@ -445,7 +445,36 @@ MP_DEFINE_CONST_FUN_OBJ_1(mp_builtin_repr_obj, mp_builtin_repr);
 STATIC mp_obj_t mp_builtin_round(size_t n_args, const mp_obj_t *args) {
     mp_obj_t o_in = args[0];
     if (MP_OBJ_IS_INT(o_in)) {
-        return o_in;
+        if (n_args <= 1) {
+            return o_in;
+        }
+
+        #if !MICROPY_PY_BUILTINS_ROUND_INT
+        mp_raise_NotImplementedError(NULL);
+        #else
+        mp_int_t num_dig = mp_obj_get_int(args[1]);
+        if (num_dig >= 0) {
+            return o_in;
+        }
+
+        mp_obj_t mult = mp_binary_op(MP_BINARY_OP_POWER, MP_OBJ_NEW_SMALL_INT(10), MP_OBJ_NEW_SMALL_INT(-num_dig));
+        mp_obj_t half_mult =  mp_binary_op(MP_BINARY_OP_FLOOR_DIVIDE, mult, MP_OBJ_NEW_SMALL_INT(2));
+        mp_obj_t modulo = mp_binary_op(MP_BINARY_OP_MODULO, o_in, mult);
+        mp_obj_t rounded = mp_binary_op(MP_BINARY_OP_SUBTRACT, o_in, modulo);
+        if (mp_obj_is_true(mp_binary_op(MP_BINARY_OP_MORE, half_mult, modulo))) {
+            return rounded;
+        } else if (mp_obj_is_true(mp_binary_op(MP_BINARY_OP_MORE, modulo, half_mult))) {
+            return mp_binary_op(MP_BINARY_OP_ADD, rounded, mult);
+        } else {
+            // round to even number
+            mp_obj_t floor = mp_binary_op(MP_BINARY_OP_FLOOR_DIVIDE, o_in, mult);
+            if (mp_obj_is_true(mp_binary_op(MP_BINARY_OP_AND, floor, MP_OBJ_NEW_SMALL_INT(1)))) {
+                return mp_binary_op(MP_BINARY_OP_ADD, rounded, mult);
+            } else {
+                return rounded;
+            }
+        }
+        #endif
     }
 #if MICROPY_PY_BUILTINS_FLOAT
     mp_float_t val = mp_obj_get_float(o_in);
diff --git a/py/mpconfig.h b/py/mpconfig.h
@@ -802,6 +802,11 @@ typedef double mp_float_t;
 #define MICROPY_PY_BUILTINS_RANGE_BINOP (0)
 #endif
 
+// Whether to support rounding of integers (incl bignum); eg round(123,-1)=120
+#ifndef MICROPY_PY_BUILTINS_ROUND_INT
+#define MICROPY_PY_BUILTINS_ROUND_INT (0)
+#endif
+
 // Whether to support timeout exceptions (like socket.timeout)
 #ifndef MICROPY_PY_BUILTINS_TIMEOUTERROR
 #define MICROPY_PY_BUILTINS_TIMEOUTERROR (0)
diff --git a/tests/basics/builtin_round_int.py b/tests/basics/builtin_round_int.py
@@ -0,0 +1,18 @@
+# test round() with integer values and second arg
+
+# rounding integers is an optional feature so test for it
+try:
+    round(1, -1)
+except NotImplementedError:
+    print('SKIP')
+    raise SystemExit
+
+tests = [
+    (1, False), (1, True),
+    (124, -1), (125, -1), (126, -1),
+    (5, -1), (15, -1), (25, -1),
+    (12345, 0), (12345, -1), (12345, 1),
+    (-1234, 0), (-1234, -1), (-1234, 1),
+]
+for t in tests:
+    print(round(*t))
diff --git a/tests/basics/builtin_round_intbig.py b/tests/basics/builtin_round_intbig.py
@@ -0,0 +1,17 @@
+# test round() with large integer values and second arg
+
+# rounding integers is an optional feature so test for it
+try:
+    round(1, -1)
+except NotImplementedError:
+    print('SKIP')
+    raise SystemExit
+
+i = 2**70
+
+tests = [
+    (i, 0), (i, -1), (i, -10), (i, 1),
+    (-i, 0), (-i, -1), (-i, -10), (-i, 1),
+]
+for t in tests:
+    print(round(*t))
