gh-158451: Reuse PyUnicodeWriter_WriteStr() in PyUnicode_Format() (#158467)

Remove special case for int in PyUnicode_Format(). Instead, reuse
PyUnicodeWriter_WriteStr() and PyUnicodeWriter_WriteRepr() in
PyUnicode_Format().

Add PyUnicode_Format() tests on integer formats.

Author: vstinner

Lib/test/test_format.py

@@ -240,6 +240,8 @@ def test_common_format(self):
         testcommon("%d", 42, "42")
         testcommon("%d", -42, "-42")
         testcommon("%d", 42.0, "42")
+        testcommon("%#d", 42, "42")
+        testcommon("%#d", -42, "-42")
         testcommon("%#x", 1, "0x1")
         testcommon("%#X", 1, "0X1")
         testcommon("%#o", 1, "0o1")
@@ -250,8 +252,12 @@ def test_common_format(self):
         testcommon("%#X", 0, "0X0")
         testcommon("%x", 0x42, "42")
         testcommon("%x", -0x42, "-42")
+        testcommon("%#x", 0x42, "0x42")
+        testcommon("%#x", -0x42, "-0x42")
         testcommon("%o", 0o42, "42")
         testcommon("%o", -0o42, "-42")
+        testcommon("%#o", 0o42, "0o42")
+        testcommon("%#o", -0o42, "-0o42")
         # alternate float formatting
         testcommon('%g', 1.1, '1.1')
         testcommon('%#g', 1.1, '1.10000')
@@ -344,12 +350,21 @@ def test_common_format(self):
                         "format argument 1: %g requires a real number, not str")
 
     def test_str_format(self):
+        testformat("%s", "abc", "abc")
         testformat("%r", "\u0378", "'\\u0378'")  # non printable
         testformat("%a", "\u0378", "'\\u0378'")  # non printable
         testformat("%r", "\u0374", "'\u0374'")   # printable
         testformat("%a", "\u0374", "'\\u0374'")  # printable
         testformat('%(x)r', {'x': 1}, '1')
 
+        # Some small ints
+        for fmt in ('s', 'r', 'a'):
+            with self.subTest(fmt=fmt):
+                testformat("%" + fmt, 42, "42")
+                testformat("%#" + fmt, 42, "42")
+                testformat("%" + fmt, -42, "-42")
+                testformat("%#" + fmt, -42, "-42")
+
         # Test exception for unknown format characters, etc.
         if verbose:
             print('Testing exceptions')

Objects/unicode_format.c

@@ -571,6 +571,7 @@ unicode_format_arg_parse(struct unicode_formatter_t *ctx,
             arg->ch = FORMAT_READ(ctx);
             ctx->fmtpos++;
         }
+        assert(arg->width >= 0);
     }
     else if (arg->ch >= '0' && arg->ch <= '9') {
         arg->width = arg->ch - '0';
@@ -590,6 +591,7 @@ unicode_format_arg_parse(struct unicode_formatter_t *ctx,
             }
             arg->width = arg->width*10 + (arg->ch - '0');
         }
+        assert(arg->width >= 0);
     }
 
     /* Parse precision. Example: "%.3f" => prec=3 */
@@ -645,6 +647,7 @@ unicode_format_arg_parse(struct unicode_formatter_t *ctx,
                 arg->prec = arg->prec*10 + (arg->ch - '0');
             }
         }
+        assert(arg->prec >= 0);
     }
 
     /* Ignore "h", "l" and "L" format prefix (ex: "%hi" or "%ls") */
@@ -701,24 +704,27 @@ unicode_format_arg_format(struct unicode_formatter_t *ctx,
     case 's':
     case 'r':
     case 'a':
-        if (PyLong_CheckExact(v) && arg->width == -1 && arg->prec == -1) {
-            /* Fast path */
-            if (_PyLong_FormatWriter(writer, v, 10, arg->flags & F_ALT) == -1)
-                return -1;
-            return 1;
+        if (arg->width < 0 && arg->prec < 0) {
+            if (arg->ch == 's') {
+                if (PyUnicodeWriter_WriteStr((PyUnicodeWriter*)writer, v) < 0) {
+                    return -1;
+                }
+                return 1;
+            }
+            else if (arg->ch == 'r') {
+                if (PyUnicodeWriter_WriteRepr((PyUnicodeWriter*)writer, v) < 0) {
+                    return -1;
+                }
+                return 1;
+            }
         }
 
-        if (PyUnicode_CheckExact(v) && arg->ch == 's') {
-            *p_str = Py_NewRef(v);
-        }
-        else {
-            if (arg->ch == 's')
-                *p_str = PyObject_Str(v);
-            else if (arg->ch == 'r')
-                *p_str = PyObject_Repr(v);
-            else
-                *p_str = PyObject_ASCII(v);
-        }
+        if (arg->ch == 's')
+            *p_str = PyObject_Str(v);
+        else if (arg->ch == 'r')
+            *p_str = PyObject_Repr(v);
+        else
+            *p_str = PyObject_ASCII(v);
         break;
 
     case 'i':

Objects/unicode_writer.c

@@ -285,45 +285,47 @@ _PyUnicodeWriter_WriteStr(_PyUnicodeWriter *writer, PyObject *str)
 
 
 int
-PyUnicodeWriter_WriteStr(PyUnicodeWriter *writer, PyObject *obj)
+PyUnicodeWriter_WriteStr(PyUnicodeWriter *pub_writer, PyObject *obj)
 {
+    _PyUnicodeWriter *writer = (_PyUnicodeWriter*)pub_writer;
     PyTypeObject *type = Py_TYPE(obj);
     if (type == &PyUnicode_Type) {
-        return _PyUnicodeWriter_WriteStr((_PyUnicodeWriter*)writer, obj);
+        return _PyUnicodeWriter_WriteStr(writer, obj);
     }
 
     if (type == &PyLong_Type) {
-        return _PyLong_FormatWriter((_PyUnicodeWriter*)writer, obj, 10, 0);
+        return _PyLong_FormatWriter(writer, obj, 10, 0);
     }
 
     PyObject *str = PyObject_Str(obj);
     if (str == NULL) {
         return -1;
     }
 
-    int res = _PyUnicodeWriter_WriteStr((_PyUnicodeWriter*)writer, str);
+    int res = _PyUnicodeWriter_WriteStr(writer, str);
     Py_DECREF(str);
     return res;
 }
 
 
 int
-PyUnicodeWriter_WriteRepr(PyUnicodeWriter *writer, PyObject *obj)
+PyUnicodeWriter_WriteRepr(PyUnicodeWriter *pub_writer, PyObject *obj)
 {
+    _PyUnicodeWriter *writer = (_PyUnicodeWriter*)pub_writer;
     if (obj == NULL) {
-        return _PyUnicodeWriter_WriteASCIIString((_PyUnicodeWriter*)writer, "<NULL>", 6);
+        return _PyUnicodeWriter_WriteASCIIString(writer, "<NULL>", 6);
     }
 
     if (Py_TYPE(obj) == &PyLong_Type) {
-        return _PyLong_FormatWriter((_PyUnicodeWriter*)writer, obj, 10, 0);
+        return _PyLong_FormatWriter(writer, obj, 10, 0);
     }
 
     PyObject *repr = PyObject_Repr(obj);
     if (repr == NULL) {
         return -1;
     }
 
-    int res = _PyUnicodeWriter_WriteStr((_PyUnicodeWriter*)writer, repr);
+    int res = _PyUnicodeWriter_WriteStr(writer, repr);
     Py_DECREF(repr);
     return res;
 }