From 643caba8cee897f22003a2bc01a0f58141811481 Mon Sep 17 00:00:00 2001
From: Jeff Epler <jepler@gmail.com>
Date: Thu, 25 Jan 2024 09:09:06 -0600
Subject: [PATCH] core: Add support for the :_b format specifier.

This groups non-decimal values by fours, bbb_bbbb_bbbb.

The SHOW_COMMA (now SHOW_SEP) flag is overloaded for both
decimal and binary contexts; use of incorrect ":,b" is not
diagnosed, and for ":_d" the separator is still printed as ","
not "_".

Signed-off-by: Jeff Epler <jepler@gmail.com>
---
 py/mpprint.c                  | 4 ++--
 py/mpprint.h                  | 2 +-
 py/mpz.c                      | 4 +++-
 py/objint.c                   | 5 +++--
 py/objstr.c                   | 6 +++++-
 tests/basics/string_format.py | 8 ++++++++
 6 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/py/mpprint.c b/py/mpprint.c
index 3218bd2f4d..193950ffb3 100644
--- a/py/mpprint.c
+++ b/py/mpprint.c
@@ -249,8 +249,8 @@ int mp_print_mp_int(const mp_print_t *print, mp_obj_t x, int base, int base_char
     prefix = prefix_buf;
 
     char comma = '\0';
-    if (flags & PF_FLAG_SHOW_COMMA) {
-        comma = ',';
+    if (flags & PF_FLAG_SHOW_SEP) {
+        comma = base == 10 ? ',' : '_';
     }
 
     // The size of this buffer is rather arbitrary. If it's not large
diff --git a/py/mpprint.h b/py/mpprint.h
index 8383ea8579..bdb5400ef9 100644
--- a/py/mpprint.h
+++ b/py/mpprint.h
@@ -33,7 +33,7 @@
 #define PF_FLAG_SPACE_SIGN        (0x004)
 #define PF_FLAG_NO_TRAILZ         (0x008)
 #define PF_FLAG_SHOW_PREFIX       (0x010)
-#define PF_FLAG_SHOW_COMMA        (0x020)
+#define PF_FLAG_SHOW_SEP          (0x020)
 #define PF_FLAG_PAD_AFTER_SIGN    (0x040)
 #define PF_FLAG_CENTER_ADJUST     (0x080)
 #define PF_FLAG_ADD_PERCENT       (0x100)
diff --git a/py/mpz.c b/py/mpz.c
index b61997e2fd..bf4d3b902b 100644
--- a/py/mpz.c
+++ b/py/mpz.c
@@ -1666,6 +1666,8 @@ size_t mpz_as_str_inpl(const mpz_t *i, unsigned int base, const char *prefix, ch
 
     size_t ilen = i->len;
 
+    int n_comma = (base == 10) ? 3 : 4;
+
     char *s = str;
     if (ilen == 0) {
         if (prefix) {
@@ -1711,7 +1713,7 @@ size_t mpz_as_str_inpl(const mpz_t *i, unsigned int base, const char *prefix, ch
                 break;
             }
         }
-        if (comma && (s - last_comma) == 3) {
+        if (comma && (s - last_comma) == n_comma) {
             *s++ = comma;
             last_comma = s;
         }
diff --git a/py/objint.c b/py/objint.c
index be5f4653a7..dcb4eeed3c 100644
--- a/py/objint.c
+++ b/py/objint.c
@@ -209,7 +209,7 @@ STATIC const uint8_t log_base2_floor[] = {
 size_t mp_int_format_size(size_t num_bits, int base, const char *prefix, char comma) {
     assert(2 <= base && base <= 16);
     size_t num_digits = num_bits / log_base2_floor[base - 1] + 1;
-    size_t num_commas = comma ? num_digits / 3 : 0;
+    size_t num_commas = comma ? (base == 10 ? num_digits / 3 : num_digits / 4): 0;
     size_t prefix_len = prefix ? strlen(prefix) : 0;
     return num_digits + num_commas + prefix_len + 2; // +1 for sign, +1 for null byte
 }
@@ -251,6 +251,7 @@ char *mp_obj_int_formatted(char **buf, size_t *buf_size, size_t *fmt_size, mp_co
         sign = '-';
     }
 
+    int n_comma = (base == 10) ? 3 : 4;
     size_t needed_size = mp_int_format_size(sizeof(fmt_int_t) * 8, base, prefix, comma);
     if (needed_size > *buf_size) {
         *buf = m_new(char, needed_size);
@@ -275,7 +276,7 @@ char *mp_obj_int_formatted(char **buf, size_t *buf_size, size_t *fmt_size, mp_co
                 c += '0';
             }
             *(--b) = c;
-            if (comma && num != 0 && b > str && (last_comma - b) == 3) {
+            if (comma && num != 0 && b > str && (last_comma - b) == n_comma) {
                 *(--b) = comma;
                 last_comma = b;
             }
diff --git a/py/objstr.c b/py/objstr.c
index 5dfe94ac4f..74b476bc5c 100644
--- a/py/objstr.c
+++ b/py/objstr.c
@@ -1213,7 +1213,11 @@ STATIC vstr_t mp_obj_str_format_helper(const char *str, const char *top, int *ar
             }
             s = str_to_int(s, stop, &width);
             if (*s == ',') {
-                flags |= PF_FLAG_SHOW_COMMA;
+                flags |= PF_FLAG_SHOW_SEP;
+                s++;
+            }
+            if (*s == '_') {
+                flags |= PF_FLAG_SHOW_SEP;
                 s++;
             }
             if (*s == '.') {
diff --git a/tests/basics/string_format.py b/tests/basics/string_format.py
index e8600f5836..6720b828ff 100644
--- a/tests/basics/string_format.py
+++ b/tests/basics/string_format.py
@@ -22,7 +22,15 @@ test("{:4o}", 123)
 test("{:4x}", 123)
 test("{:4X}", 123)
 
+test("{:4,d}", 1)
+test("{:4_o}", 1)
+test("{:4_b}", 1)
+test("{:4_x}", 1)
+
 test("{:4,d}", 12345678)
+test("{:4_o}", 12345678)
+test("{:4_b}", 12345678)
+test("{:4_x}", 12345678)
 
 test("{:#4b}", 10)
 test("{:#4o}", 123)