py: Fix builtin abs so it works for bools and bignum.

5 files changed, 61 insertions, 8 deletions

diff --git a/py/modbuiltins.c b/py/modbuiltins.c
index 31ed8a7159..78fa4afc33 100644
--- a/py/modbuiltins.c
+++ b/py/modbuiltins.c
@@ -29,6 +29,7 @@
 
 #include "py/nlr.h"
 #include "py/smallint.h"
+#include "py/objint.h"
 #include "py/objstr.h"
 #include "py/runtime0.h"
 #include "py/runtime.h"
@@ -86,12 +87,8 @@ STATIC mp_obj_t mp_builtin___build_class__(mp_uint_t n_args, const mp_obj_t *arg
 MP_DEFINE_CONST_FUN_OBJ_VAR(mp_builtin___build_class___obj, 2, mp_builtin___build_class__);
 
 STATIC mp_obj_t mp_builtin_abs(mp_obj_t o_in) {
-    if (MP_OBJ_IS_SMALL_INT(o_in)) {
-        mp_int_t val = MP_OBJ_SMALL_INT_VALUE(o_in);
-        if (val < 0) {
-            val = -val;
-        }
-        return MP_OBJ_NEW_SMALL_INT(val);
+    if (0) {
+        // dummy
 #if MICROPY_PY_BUILTINS_FLOAT
     } else if (MP_OBJ_IS_TYPE(o_in, &mp_type_float)) {
         mp_float_t value = mp_obj_float_get(o_in);
@@ -109,8 +106,8 @@ STATIC mp_obj_t mp_builtin_abs(mp_obj_t o_in) {
 #endif
 #endif
     } else {
-        assert(0);
-        return mp_const_none;
+        // this will raise a TypeError if the argument is not integral
+        return mp_obj_int_abs(o_in);
     }
 }
 MP_DEFINE_CONST_FUN_OBJ_1(mp_builtin_abs_obj, mp_builtin_abs);
diff --git a/py/objint.c b/py/objint.c
index 0fb2a38262..26e1cc048f 100644
--- a/py/objint.c
+++ b/py/objint.c
@@ -266,6 +266,16 @@ bool mp_obj_int_is_positive(mp_obj_t self_in) {
     return mp_obj_get_int(self_in) >= 0;
 }
 
+// This must handle int and bool types, and must raise a
+// TypeError if the argument is not integral
+mp_obj_t mp_obj_int_abs(mp_obj_t self_in) {
+    mp_int_t val = mp_obj_get_int(self_in);
+    if (val < 0) {
+        val = -val;
+    }
+    return MP_OBJ_NEW_SMALL_INT(val);
+}
+
 // This is called for operations on SMALL_INT that are not handled by mp_unary_op
 mp_obj_t mp_obj_int_unary_op(mp_uint_t op, mp_obj_t o_in) {
     return MP_OBJ_NULL; // op not supported
diff --git a/py/objint.h b/py/objint.h
index c995f2632a..63cd91a698 100644
--- a/py/objint.h
+++ b/py/objint.h
@@ -57,6 +57,7 @@ char *mp_obj_int_formatted_impl(char **buf, mp_uint_t *buf_size, mp_uint_t *fmt_
                                 int base, const char *prefix, char base_char, char comma);
 mp_int_t mp_obj_int_hash(mp_obj_t self_in);
 bool mp_obj_int_is_positive(mp_obj_t self_in);
+mp_obj_t mp_obj_int_abs(mp_obj_t self_in);
 mp_obj_t mp_obj_int_unary_op(mp_uint_t op, mp_obj_t o_in);
 mp_obj_t mp_obj_int_binary_op(mp_uint_t op, mp_obj_t lhs_in, mp_obj_t rhs_in);
 mp_obj_t mp_obj_int_binary_op_extra_cases(mp_uint_t op, mp_obj_t lhs_in, mp_obj_t rhs_in);
diff --git a/py/objint_longlong.c b/py/objint_longlong.c
index 83e1c67d6d..837889704a 100644
--- a/py/objint_longlong.c
+++ b/py/objint_longlong.c
@@ -71,6 +71,30 @@ bool mp_obj_int_is_positive(mp_obj_t self_in) {
     return self->val >= 0;
 }
 
+// This must handle int and bool types, and must raise a
+// TypeError if the argument is not integral
+mp_obj_t mp_obj_int_abs(mp_obj_t self_in) {
+    if (MP_OBJ_IS_TYPE(self_in, &mp_type_int)) {
+        mp_obj_int_t *self = self_in;
+        self = mp_obj_new_int_from_ll(self->val);
+        if (self->val < 0) {
+            // TODO could overflow long long
+            self->val = -self->val;
+        }
+        return self;
+    } else {
+        mp_int_t val = mp_obj_get_int(self_in);
+        if (val == MP_SMALL_INT_MIN) {
+            return mp_obj_new_int_from_ll(-val);
+        } else {
+            if (val < 0) {
+                val = -val;
+            }
+            return MP_OBJ_NEW_SMALL_INT(val);
+        }
+    }
+}
+
 mp_obj_t mp_obj_int_unary_op(mp_uint_t op, mp_obj_t o_in) {
     mp_obj_int_t *o = o_in;
     switch (op) {
diff --git a/py/objint_mpz.c b/py/objint_mpz.c
index c46692ec76..be917f4a74 100644
--- a/py/objint_mpz.c
+++ b/py/objint_mpz.c
@@ -112,6 +112,27 @@ bool mp_obj_int_is_positive(mp_obj_t self_in) {
     return !self->mpz.neg;
 }
 
+// This must handle int and bool types, and must raise a
+// TypeError if the argument is not integral
+mp_obj_t mp_obj_int_abs(mp_obj_t self_in) {
+    if (MP_OBJ_IS_TYPE(self_in, &mp_type_int)) {
+        mp_obj_int_t *self = self_in;
+        mp_obj_int_t *self2 = mp_obj_int_new_mpz();
+        mpz_abs_inpl(&self2->mpz, &self->mpz);
+        return self2;
+    } else {
+        mp_int_t val = mp_obj_get_int(self_in);
+        if (val == MP_SMALL_INT_MIN) {
+            return mp_obj_new_int_from_ll(-val);
+        } else {
+            if (val < 0) {
+                val = -val;
+            }
+            return MP_OBJ_NEW_SMALL_INT(val);
+        }
+    }
+}
+
 mp_obj_t mp_obj_int_unary_op(mp_uint_t op, mp_obj_t o_in) {
     mp_obj_int_t *o = o_in;
     switch (op) {