pyo3/
fmt.rs

1// TODO https://github.com/PyO3/pyo3/issues/5487
2#![allow(clippy::undocumented_unsafe_blocks)]
3
4#[allow(unused_imports, reason = "used to build docs")]
5use crate::platform::prelude::*;
6#[cfg(any(doc, all(Py_3_14, not(Py_LIMITED_API))))]
7use crate::{types::PyString, Python};
8#[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
9use {
10    crate::ffi::{
11        PyUnicodeWriter_Create, PyUnicodeWriter_Discard, PyUnicodeWriter_Finish,
12        PyUnicodeWriter_WriteChar, PyUnicodeWriter_WriteUTF8,
13    },
14    crate::ffi_ptr_ext::FfiPtrExt,
15    crate::py_result_ext::PyResultExt,
16    crate::IntoPyObject,
17    crate::{ffi, Bound, PyErr, PyResult},
18    core::fmt,
19    core::mem::ManuallyDrop,
20    core::ptr::NonNull,
21};
22
23/// This macro is analogous to Rust's [`format!`] macro, but returns a [`PyString`] instead of a [`String`].
24///
25/// # Arguments
26///
27/// The arguments are exactly like [`format!`], but with `py` (a [`Python`] token) as the first argument:
28///
29/// # Interning Advantage
30///
31/// If the format string is a static string and all arguments are constant at compile time,
32/// this macro will intern the string in Python, offering better performance and memory usage
33/// compared to [`PyString::from_fmt`].
34///
35/// ```rust
36/// # use pyo3::{py_format, Python, types::PyString, Bound};
37/// Python::attach(|py| {
38///     let py_string: Bound<'_, PyString> = py_format!(py, "{} {}", "hello", "world").unwrap();
39///     assert_eq!(py_string.to_string(), "hello world");
40/// });
41/// ```
42#[macro_export]
43macro_rules! py_format {
44    ($py: expr, $($arg:tt)*) => {{
45        if let Some(static_string) = format_args!($($arg)*).as_str() {
46            static INTERNED: $crate::sync::PyOnceLock<$crate::Py<$crate::types::PyString>> = $crate::sync::PyOnceLock::new();
47            Ok($crate::Bound::clone(
48                INTERNED
49                .get_or_init($py, || $crate::types::PyString::intern($py, static_string).unbind())
50                .bind($py)
51            ))
52        } else {
53            $crate::types::PyString::from_fmt($py, format_args!($($arg)*))
54        }
55    }}
56}
57
58#[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
59/// The `PyUnicodeWriter` is a utility for efficiently constructing Python strings
60pub(crate) struct PyUnicodeWriter<'py> {
61    python: Python<'py>,
62    writer: NonNull<ffi::PyUnicodeWriter>,
63    last_error: Option<PyErr>,
64}
65
66#[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
67impl<'py> PyUnicodeWriter<'py> {
68    /// Creates a new `PyUnicodeWriter`.
69    pub fn new(py: Python<'py>) -> PyResult<Self> {
70        Self::with_capacity(py, 0)
71    }
72
73    /// Creates a new `PyUnicodeWriter` with the specified initial capacity.
74    #[inline]
75    pub fn with_capacity(py: Python<'py>, capacity: usize) -> PyResult<Self> {
76        match NonNull::new(unsafe { PyUnicodeWriter_Create(capacity.try_into()?) }) {
77            Some(ptr) => Ok(PyUnicodeWriter {
78                python: py,
79                writer: ptr,
80                last_error: None,
81            }),
82            None => Err(PyErr::fetch(py)),
83        }
84    }
85
86    /// Consumes the `PyUnicodeWriter` and returns a `Bound<PyString>` containing the constructed string.
87    #[inline]
88    pub fn into_py_string(mut self) -> PyResult<Bound<'py, PyString>> {
89        let py = self.python;
90        if let Some(error) = self.take_error() {
91            Err(error)
92        } else {
93            unsafe {
94                PyUnicodeWriter_Finish(ManuallyDrop::new(self).as_ptr())
95                    .assume_owned_or_err(py)
96                    .cast_into_unchecked()
97            }
98        }
99    }
100
101    /// When fmt::Write returned an error, this function can be used to retrieve the last error that occurred.
102    #[inline]
103    pub fn take_error(&mut self) -> Option<PyErr> {
104        self.last_error.take()
105    }
106
107    #[inline]
108    fn as_ptr(&self) -> *mut ffi::PyUnicodeWriter {
109        self.writer.as_ptr()
110    }
111
112    #[inline]
113    fn set_error(&mut self) {
114        self.last_error = Some(PyErr::fetch(self.python));
115    }
116}
117
118#[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
119impl fmt::Write for PyUnicodeWriter<'_> {
120    #[inline]
121    fn write_str(&mut self, s: &str) -> fmt::Result {
122        let result = unsafe {
123            PyUnicodeWriter_WriteUTF8(self.as_ptr(), s.as_ptr().cast(), s.len() as isize)
124        };
125        if result < 0 {
126            self.set_error();
127            Err(fmt::Error)
128        } else {
129            Ok(())
130        }
131    }
132
133    #[inline]
134    fn write_char(&mut self, c: char) -> fmt::Result {
135        let result = unsafe { PyUnicodeWriter_WriteChar(self.as_ptr(), c.into()) };
136        if result < 0 {
137            self.set_error();
138            Err(fmt::Error)
139        } else {
140            Ok(())
141        }
142    }
143}
144
145#[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
146impl Drop for PyUnicodeWriter<'_> {
147    #[inline]
148    fn drop(&mut self) {
149        unsafe {
150            PyUnicodeWriter_Discard(self.as_ptr());
151        }
152    }
153}
154
155#[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
156impl<'py> IntoPyObject<'py> for PyUnicodeWriter<'py> {
157    type Target = PyString;
158    type Output = Bound<'py, Self::Target>;
159    type Error = PyErr;
160
161    #[inline]
162    fn into_pyobject(self, _py: Python<'py>) -> PyResult<Bound<'py, PyString>> {
163        self.into_py_string()
164    }
165}
166
167#[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
168impl<'py> TryInto<Bound<'py, PyString>> for PyUnicodeWriter<'py> {
169    type Error = PyErr;
170
171    #[inline]
172    fn try_into(self) -> PyResult<Bound<'py, PyString>> {
173        self.into_py_string()
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    #[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
180    use super::*;
181    use crate::types::PyStringMethods;
182    use crate::{IntoPyObject, Python};
183
184    #[test]
185    #[allow(clippy::write_literal)]
186    #[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
187    fn unicode_writer_test() {
188        use core::fmt::Write;
189        Python::attach(|py| {
190            let mut writer = PyUnicodeWriter::new(py).unwrap();
191            write!(writer, "Hello {}!", "world").unwrap();
192            writer.write_char('😎').unwrap();
193            let result = writer.into_py_string().unwrap();
194            assert_eq!(result.to_string(), "Hello world!😎");
195        });
196    }
197
198    #[test]
199    #[allow(clippy::write_literal)]
200    #[cfg(all(Py_3_14, not(Py_LIMITED_API)))]
201    fn unicode_writer_with_capacity() {
202        use core::fmt::Write;
203        Python::attach(|py| {
204            let mut writer = PyUnicodeWriter::with_capacity(py, 10).unwrap();
205            write!(writer, "Hello {}!", "world").unwrap();
206            writer.write_char('😎').unwrap();
207            let result = writer.into_py_string().unwrap();
208            assert_eq!(result.to_string(), "Hello world!😎");
209        });
210    }
211
212    #[test]
213    fn test_pystring_from_fmt() {
214        Python::attach(|py| {
215            py_format!(py, "Hello {}!", "world").unwrap();
216        });
217    }
218
219    #[test]
220    fn test_complex_format() {
221        Python::attach(|py| {
222            let complex_value = (42, "foo", [0; 0]).into_pyobject(py).unwrap();
223            let py_string = py_format!(py, "This is some complex value: {complex_value}").unwrap();
224            let actual = py_string.to_cow().unwrap();
225            let expected = "This is some complex value: (42, 'foo', [])";
226            assert_eq!(actual, expected);
227        });
228    }
229}
pyo3/fmt.rs

pyo3/
fmt.rs
