Add PyString::from_fmt using new PyUnicodeWriter

bschoenmaeckers · bschoenmaeckers · commit bfb72fb8e93a · 2025-06-19T15:08:53.000+02:00
diff --git a/pyo3-ffi/src/compat/py_3_14.rs b/pyo3-ffi/src/compat/py_3_14.rs
@@ -24,3 +24,83 @@ compat_function!(
         }
     }
 );
+
+compat_function!(
+    originally_defined_for(all(Py_3_14, not(Py_LIMITED_API)));
+
+    pub unsafe fn PyUnicodeWriter_Create(length: crate::Py_ssize_t) -> *mut crate::PyUnicodeWriter {
+        if length < 0 {
+            crate::PyErr_SetString(
+                crate::PyExc_ValueError,
+                c_str!("length must be positive").as_ptr(),
+            );
+            return std::ptr::null_mut();
+        }
+
+        let size = std::mem::size_of::<crate::_PyUnicodeWriter>();
+        let writer: *mut crate::_PyUnicodeWriter = crate::PyMem_Malloc(size).cast();
+        crate::_PyUnicodeWriter_Init(writer);
+        if crate::_PyUnicodeWriter_Prepare(writer, length, 127) < 0 {
+            PyUnicodeWriter_Discard(writer.cast());
+            return std::ptr::null_mut();
+        }
+        (*writer).overallocate = 1;
+        writer.cast()
+    }
+);
+
+compat_function!(
+    originally_defined_for(all(Py_3_14, not(Py_LIMITED_API)));
+
+    pub unsafe fn PyUnicodeWriter_Finish(writer: *mut crate::PyUnicodeWriter) -> *mut crate::PyObject {
+        let str = crate::_PyUnicodeWriter_Finish(writer.cast());
+        crate::PyMem_Free(writer.cast());
+        str
+    }
+);
+
+compat_function!(
+    originally_defined_for(all(Py_3_14, not(Py_LIMITED_API)));
+
+    pub unsafe fn PyUnicodeWriter_Discard(writer: *mut crate::PyUnicodeWriter) -> () {
+        crate::_PyUnicodeWriter_Dealloc(writer.cast());
+        crate::PyMem_Free(writer.cast())
+    }
+);
+
+compat_function!(
+    originally_defined_for(all(Py_3_14, not(Py_LIMITED_API)));
+
+    pub unsafe fn PyUnicodeWriter_WriteChar(writer: *mut crate::PyUnicodeWriter, ch: crate::Py_UCS4) -> std::os::raw::c_int {
+        if ch > 0x10ffff {
+            crate::PyErr_SetString(
+                crate::PyExc_ValueError,
+                c_str!("character must be in range(0x110000)").as_ptr(),
+            );
+            return -1;
+        }
+
+        crate::_PyUnicodeWriter_WriteChar(writer.cast(), ch)
+    }
+);
+
+compat_function!(
+    originally_defined_for(all(Py_3_14, not(Py_LIMITED_API)));
+
+    pub unsafe fn PyUnicodeWriter_WriteUTF8(writer: *mut crate::PyUnicodeWriter,str: *const std::os::raw::c_char, size: crate::Py_ssize_t) -> std::os::raw::c_int {
+        let size = if size < 0 {
+            libc::strlen(str) as isize
+        } else {
+            size
+        };
+
+        let py_str = crate::PyUnicode_FromStringAndSize(str, size);
+        if py_str.is_null() {
+            return -1;
+        }
+
+        let result = crate::_PyUnicodeWriter_WriteStr(writer.cast(), py_str);
+        crate::Py_DECREF(py_str);
+        result
+    }
+);
diff --git a/pyo3-ffi/src/cpython/unicodeobject.rs b/pyo3-ffi/src/cpython/unicodeobject.rs
@@ -686,6 +686,77 @@ extern "C" {
     // skipped PyUnicode_GetMax
 }
 
+opaque_struct!(pub PyUnicodeWriter);
+
+#[cfg(not(Py_3_14))]
+#[repr(C)]
+pub(crate) struct _PyUnicodeWriter {
+    buffer: *mut PyObject,
+    data: *mut c_void,
+    kind: c_int,
+    pub(crate) maxchar: Py_UCS4,
+    pub(crate) size: Py_ssize_t,
+    pub(crate) pos: Py_ssize_t,
+    min_length: Py_ssize_t,
+    min_char: Py_UCS4,
+    pub(crate) overallocate: c_char,
+    readonly: c_char,
+}
+
+extern "C" {
+    #[cfg(Py_3_14)]
+    pub fn PyUnicodeWriter_Create(length: Py_ssize_t) -> *mut PyUnicodeWriter;
+    #[cfg(Py_3_14)]
+    pub fn PyUnicodeWriter_Finish(writer: *mut PyUnicodeWriter) -> *mut PyObject;
+    #[cfg(not(Py_3_14))]
+    pub(crate) fn _PyUnicodeWriter_Finish(writer: *mut _PyUnicodeWriter) -> *mut PyObject;
+    #[cfg(Py_3_14)]
+    pub fn PyUnicodeWriter_Discard(writer: *mut PyUnicodeWriter);
+    #[cfg(not(Py_3_14))]
+    pub(crate) fn _PyUnicodeWriter_Dealloc(writer: *mut _PyUnicodeWriter);
+    #[cfg(not(Py_3_14))]
+    pub(crate) fn _PyUnicodeWriter_Init(writer: *mut _PyUnicodeWriter);
+    #[cfg(not(Py_3_14))]
+    pub(crate) fn _PyUnicodeWriter_PrepareInternal(
+        writer: *mut _PyUnicodeWriter,
+        length: Py_ssize_t,
+        maxchars: Py_UCS4,
+    ) -> c_int;
+    #[cfg(Py_3_14)]
+    pub fn PyUnicodeWriter_WriteChar(writer: *mut PyUnicodeWriter, ch: Py_UCS4) -> c_int;
+    #[cfg(not(Py_3_14))]
+    pub(crate) fn _PyUnicodeWriter_WriteChar(writer: *mut _PyUnicodeWriter, ch: Py_UCS4) -> c_int;
+    #[cfg(not(Py_3_14))]
+    pub(crate) fn _PyUnicodeWriter_WriteStr(
+        writer: *mut _PyUnicodeWriter,
+        str: *mut PyObject,
+    ) -> c_int;
+    #[cfg(Py_3_14)]
+    pub fn PyUnicodeWriter_WriteUTF8(
+        writer: *mut PyUnicodeWriter,
+        str: *const c_char,
+        size: Py_ssize_t,
+    ) -> c_int;
+}
+
+#[cfg(not(Py_3_14))]
+#[inline(always)]
+pub(crate) unsafe fn _PyUnicodeWriter_Prepare(
+    writer: *mut _PyUnicodeWriter,
+    length: Py_ssize_t,
+    maxchars: Py_UCS4,
+) -> c_int {
+    if maxchars <= (*writer).maxchar && length <= (*writer).size - (*writer).pos {
+        return 0;
+    }
+
+    if length == 0 {
+        return 0;
+    }
+
+    _PyUnicodeWriter_PrepareInternal(writer, length, maxchars)
+}
+
 // skipped _PyUnicodeWriter
 // skipped _PyUnicodeWriter_Init
 // skipped _PyUnicodeWriter_Prepare
diff --git a/src/fmt.rs b/src/fmt.rs
@@ -0,0 +1,137 @@
+//! This module provides the `PyUnicodeWriter` struct, which is a utility for efficiently
+//! constructing Python strings using Rust's `fmt::Write` trait.
+//! It allows for incremental string construction, without the need for repeated allocations, and
+//! is particularly useful for building strings in a performance-sensitive context.
+use crate::ffi::compat::{
+    PyUnicodeWriter_Create, PyUnicodeWriter_Discard, PyUnicodeWriter_Finish,
+    PyUnicodeWriter_WriteChar, PyUnicodeWriter_WriteUTF8,
+};
+use crate::ffi_ptr_ext::FfiPtrExt;
+use crate::impl_::callback::WrappingCastTo;
+use crate::types::{PyAnyMethods, PyString};
+use crate::{ffi, Bound, PyErr, PyResult, Python};
+use std::ptr::NonNull;
+use std::{fmt, mem};
+
+/// The `PyUnicodeWriter` is a utility for efficiently constructing Python strings
+pub struct PyUnicodeWriter {
+    writer: NonNull<ffi::PyUnicodeWriter>,
+    last_error: Option<PyErr>,
+}
+
+impl PyUnicodeWriter {
+    /// Creates a new `PyUnicodeWriter`.
+    pub fn new(py: Python<'_>) -> PyResult<Self> {
+        Self::with_capacity(py, 0)
+    }
+
+    /// Creates a new `PyUnicodeWriter` with the specified initial capacity.
+    pub fn with_capacity(py: Python<'_>, capacity: usize) -> PyResult<Self> {
+        match NonNull::new(unsafe { PyUnicodeWriter_Create(capacity.wrapping_cast()) }) {
+            Some(ptr) => Ok(PyUnicodeWriter {
+                writer: ptr,
+                last_error: None,
+            }),
+            None => Err(PyErr::fetch(py)),
+        }
+    }
+
+    /// Consumes the `PyUnicodeWriter` and returns a `Bound<PyString>` containing the constructed string.
+    pub fn into_py_string(self, py: Python<'_>) -> PyResult<Bound<'_, PyString>> {
+        let writer_ptr = self.as_ptr();
+        mem::forget(self);
+        Ok(unsafe {
+            PyUnicodeWriter_Finish(writer_ptr)
+                .assume_owned_or_err(py)?
+                .downcast_into_unchecked()
+        })
+    }
+
+    /// When fmt::Write returned an error, this function can be used to retrieve the last error that occurred.
+    pub fn take_error(&mut self) -> Option<PyErr> {
+        self.last_error.take()
+    }
+
+    fn as_ptr(&self) -> *mut ffi::PyUnicodeWriter {
+        self.writer.as_ptr()
+    }
+
+    fn set_error(&mut self) {
+        Python::with_gil(|py| {
+            self.last_error = Some(PyErr::fetch(py));
+        })
+    }
+}
+
+impl fmt::Write for PyUnicodeWriter {
+    fn write_str(&mut self, s: &str) -> fmt::Result {
+        let result = unsafe {
+            PyUnicodeWriter_WriteUTF8(self.as_ptr(), s.as_ptr().cast(), s.len() as isize)
+        };
+        if result < 0 {
+            self.set_error();
+            Err(fmt::Error)
+        } else {
+            Ok(())
+        }
+    }
+
+    fn write_char(&mut self, c: char) -> fmt::Result {
+        let result = unsafe { PyUnicodeWriter_WriteChar(self.as_ptr(), c as u32) };
+        if result < 0 {
+            self.set_error();
+            Err(fmt::Error)
+        } else {
+            Ok(())
+        }
+    }
+}
+
+impl Drop for PyUnicodeWriter {
+    fn drop(&mut self) {
+        unsafe {
+            PyUnicodeWriter_Discard(self.as_ptr());
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::types::PyStringMethods;
+    use crate::{IntoPyObject, Python};
+    use std::fmt::Write;
+
+    #[test]
+    fn unicode_writer_test() {
+        Python::with_gil(|py| {
+            let mut writer = PyUnicodeWriter::new(py).unwrap();
+            write!(writer, "Hello {}!", "world").unwrap();
+            writer.write_char('😎').unwrap();
+            let result = writer.into_py_string(py).unwrap();
+            assert_eq!(result.to_string(), "Hello world!😎");
+        });
+    }
+
+    #[test]
+    fn test_pystring_from_fmt() {
+        Python::with_gil(|py| {
+            PyString::from_fmt(py, format_args!("Hello {}!", "world")).unwrap();
+        });
+    }
+
+    #[test]
+    fn test_complex_format() {
+        Python::with_gil(|py| {
+            let complex_value = (42, "foo", 3.14).into_pyobject(py).unwrap();
+            let py_string = PyString::from_fmt(
+                py,
+                format_args!("This is some complex value: {complex_value}"),
+            )
+            .unwrap();
+            let actual = py_string.to_cow().unwrap();
+            let expected = "This is some complex value: (42, 'foo', 3.14)";
+            assert_eq!(actual, expected);
+        });
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
@@ -421,6 +421,7 @@ pub mod coroutine;
 mod err;
 pub mod exceptions;
 pub mod ffi;
+pub mod fmt;
 mod gil;
 #[doc(hidden)]
 pub mod impl_;
diff --git a/src/types/string.rs b/src/types/string.rs
@@ -1,6 +1,7 @@
 #[cfg(not(Py_LIMITED_API))]
 use crate::exceptions::PyUnicodeDecodeError;
 use crate::ffi_ptr_ext::FfiPtrExt;
+use crate::fmt::PyUnicodeWriter;
 use crate::instance::Borrowed;
 use crate::py_result_ext::PyResultExt;
 use crate::types::any::PyAnyMethods;
@@ -9,7 +10,8 @@ use crate::types::PyBytes;
 use crate::{ffi, Bound, Py, PyAny, PyResult, Python};
 use std::borrow::Cow;
 use std::ffi::CString;
-use std::str;
+use std::fmt::Write as _;
+use std::{fmt, str};
 
 /// Represents raw data backing a Python `str`.
 ///
@@ -209,6 +211,24 @@ impl PyString {
             .downcast_into_unchecked()
         }
     }
+
+    /// Creates a Python string using a format string.
+    ///
+    /// This function is similar to [`format!`], but it returns a Python string object instead of a Rust string.
+    pub fn from_fmt<'py>(
+        py: Python<'py>,
+        args: fmt::Arguments<'_>,
+    ) -> PyResult<Bound<'py, PyString>> {
+        if let Some(static_string) = args.as_str() {
+            return Ok(PyString::new(py, static_string));
+        };
+
+        let mut writer = PyUnicodeWriter::new(py)?;
+        writer
+            .write_fmt(args)
+            .map_err(|_| writer.take_error().expect("expected error"))?;
+        writer.into_py_string(py)
+    }
 }
 
 /// Implementation of functionality for [`PyString`].
