Add string formatting support for longlong and mpz.

1 files changed, 64 insertions, 45 deletions

diff --git a/py/pfenv.c b/py/pfenv.c
index 0d6fab3c48..ab45328ad7 100644
--- a/py/pfenv.c
+++ b/py/pfenv.c
@@ -5,6 +5,8 @@
 #include "mpconfig.h"
 #include "qstr.h"
 #include "obj.h"
+#include "mpz.h"
+#include "objint.h"
 #include "pfenv.h"
 
 #if MICROPY_FLOAT_IMPL == MICROPY_FLOAT_IMPL_DOUBLE
@@ -79,71 +81,88 @@ int pfenv_print_strn(const pfenv_t *pfenv, const char *str, unsigned int len, in
 // We can use 16 characters for 32-bit and 32 characters for 64-bit
 #define INT_BUF_SIZE (sizeof(machine_int_t) * 4)
 
-int pfenv_print_int(const pfenv_t *pfenv, machine_uint_t x, int sgn, int base, int base_char, int flags, char fill, int width) {
-    char sign = 0;
-    if (sgn) {
-        if ((machine_int_t)x < 0) {
-            sign = '-';
-            x = -x;
-        } else if (flags & PF_FLAG_SHOW_SIGN) {
-            sign = '+';
-        } else if (flags & PF_FLAG_SPACE_SIGN) {
-            sign = ' ';
-        }
+int pfenv_print_int(const pfenv_t *pfenv, mp_obj_t x, int sgn, int base, int base_char, int flags, char fill, int width) {
+    if (!MP_OBJ_IS_INT(x)) {
+        // This will convert booleans to int, or raise an error for
+        // non-integer types.
+        x = MP_OBJ_NEW_SMALL_INT(mp_obj_get_int(x));
     }
 
-    char buf[INT_BUF_SIZE];
-    char *b = buf + INT_BUF_SIZE;
+    char prefix_buf[4];
+    char *prefix = prefix_buf;
 
-    if (x == 0) {
-        *(--b) = '0';
-    } else {
-        do {
-            int c = x % base;
-            x /= base;
-            if (c >= 10) {
-                c += base_char - 10;
-            } else {
-                c += '0';
-            }
-            *(--b) = c;
-        } while (b > buf && x != 0);
+    if (mp_obj_int_is_positive(x)) {
+        if (flags & PF_FLAG_SHOW_SIGN) {
+            *prefix++ = '+';
+        } else if (flags & PF_FLAG_SPACE_SIGN) {
+            *prefix++ = ' ';
+        }
     }
 
-    char prefix_char = '\0';
-
     if (flags & PF_FLAG_SHOW_PREFIX) {
         if (base == 2) {
-            prefix_char = base_char + 'b' - 'a';
+            *prefix++ = '0';
+            *prefix++ = base_char + 'b' - 'a';
         } else if (base == 8) {
-            prefix_char = base_char + 'o' - 'a';
+            *prefix++ = '0';
+            if (flags & PF_FLAG_SHOW_OCTAL_LETTER) {
+                *prefix++ = base_char + 'o' - 'a';
+            }
         } else if (base == 16) {
-            prefix_char = base_char + 'x' - 'a';
+            *prefix++ = '0';
+            *prefix++ = base_char + 'x' - 'a';
+        }
+    }
+    *prefix = '\0';
+    int prefix_len = prefix - prefix_buf;
+    prefix = prefix_buf;
+
+    char comma = '\0';
+    if (flags & PF_FLAG_SHOW_COMMA) {
+        comma = ',';
+    }
+
+    // The size of this buffer is rather arbitrary. If it's not large
+    // enough, a dynamic one will be allocated.
+    char stack_buf[sizeof(machine_int_t) * 4];
+    char *buf = stack_buf;
+    int buf_size = sizeof(stack_buf);
+    int fmt_size = 0;
+    char *str;
+
+    char sign = '\0';
+    if (flags & PF_FLAG_PAD_AFTER_SIGN) {
+        // We add the pad in this function, so since the pad goes after
+        // the sign & prefix, we format without a prefix
+        str = mp_obj_int_formatted(&buf, &buf_size, &fmt_size,
+                                   x, base, NULL, base_char, comma);
+        if (*str == '-') {
+            sign = *str++;
+            fmt_size--;
         }
+    } else {
+        str = mp_obj_int_formatted(&buf, &buf_size, &fmt_size,
+                                   x, base, prefix, base_char, comma);
     }
 
     int len = 0;
     if (flags & PF_FLAG_PAD_AFTER_SIGN) {
+        // pad after sign implies pad after prefix as well.
         if (sign) {
-            len += pfenv_print_strn(pfenv, &sign, 1, flags, fill, 1);
+            len += pfenv_print_strn(pfenv, &sign, 1, 0, 0, 1);
             width--;
         }
-        if (prefix_char) {
-            len += pfenv_print_strn(pfenv, "0", 1, flags, fill, 1);
-            len += pfenv_print_strn(pfenv, &prefix_char, 1, flags, fill, 1);
-            width -= 2;
-        }
-    } else {
-        if (prefix_char && b > &buf[1]) {
-            *(--b) = prefix_char;
-            *(--b) = '0';
-        }
-        if (sign && b > buf) {
-            *(--b) = sign;
+        if (prefix_len) {
+            len += pfenv_print_strn(pfenv, prefix, prefix_len, 0, 0, 1);
+            width -= prefix_len;
         }
     }
 
-    len += pfenv_print_strn(pfenv, b, buf + INT_BUF_SIZE - b, flags, fill, width);
+    len += pfenv_print_strn(pfenv, str, fmt_size, flags, fill, width);
+
+    if (buf != stack_buf) {
+        m_free(buf, buf_size);
+    }
     return len;
 }